{"lightning_attention.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/lightning_attention.py_gen_triton_code_194184.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "dequantize_rowwise.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/dequantize_rowwise.py_gen_triton_code_421993.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "kldiv_compute.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/kldiv_compute.py_gen_triton_code_227847.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "context_attn_bloom.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/context_attn_bloom.py_gen_triton_code_62193.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "swiglu_fwd.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/swiglu_fwd.py_gen_triton_code_401619.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "int8_matmul_kernel.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/int8_matmul_kernel.py_gen_triton_code_427291.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "attention_kernel_aligned.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/attention_kernel_aligned.py_gen_triton_code_657497.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "llama_ff_triton.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/llama_ff_triton.py_gen_triton_code_166248.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\nfrom triton import Config, autotune\n\n@autotune(\n    configs=[\n        Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=3, num_warps=8),\n        Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=3, num_warps=8),\n        Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=5, num_warps=2),\n        Config({'BLOCK_SIZE_M': 16, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=2, num_warps=4),\n        Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=3, num_warps=8),\n        Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=3, num_warps=8),\n        Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=4, num_warps=4),\n        Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8, 'SPLIT_K': 1}, num_stages=5, num_warps=2),\n    ],\n    key=['M', 'N', 'K'],\n)\n@triton.jit\ndef ff_llama(\n    a_ptr, w1_ptr, w3_ptr, \n    out_ptr, act_in_ptr, mm_1_ptr,\n    M, N, K,\n    stride_am, stride_ak,\n    stride_w1k, stride_w1n,\n    stride_w3k, stride_w3n,\n    stride_outm, stride_outn,\n    stride_inm, stride_inn,\n    stride_mmm, stride_mmn,\n    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,\n    GROUP_SIZE_M: tl.constexpr, SPLIT_K: tl.constexpr,\n):\n    \"\"\"\n    w1 and w3 are weights (linear layers)\n    F.silu(w1(x)) * w3(x)\n    \"\"\"\n    pid = tl.program_id(axis=0)\n    pid_m = pid // tl.cdiv(N, BLOCK_SIZE_N)\n    pid_n = pid % tl.cdiv(N, BLOCK_SIZE_N)\n\n    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M\n    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N\n    offs_k = tl.arange(0, BLOCK_SIZE_K)\n    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)\n    w1_ptrs = w1_ptr + (offs_k[:, None] * stride_w1k + offs_bn[None, :] * stride_w1n)\n    w3_ptrs = w3_ptr + (offs_k[:, None] * stride_w3k + offs_bn[None, :] * stride_w3n)\n    acc1 = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n    acc2 = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n\n    for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):\n        a = tl.load(a_ptrs)\n        b = tl.load(w1_ptrs)\n        acc1 += tl.dot(a, b)\n\n        c = tl.load(w3_ptrs)\n        acc2 += tl.dot(a, c)\n\n        a_ptrs += BLOCK_SIZE_K * stride_ak\n        w1_ptrs += BLOCK_SIZE_K * stride_w1k\n        w3_ptrs += BLOCK_SIZE_K * stride_w3k\n\n    accumulator = (acc1 * tl.sigmoid(acc1)) * acc2\n\n    offs_outm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)\n    offs_outn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)\n\n    out_ptrs = out_ptr + (stride_outm * offs_outm[:, None] + stride_outn * offs_outn[None, :])\n    act_in_ptrs = act_in_ptr + (stride_inm * offs_outm[:, None] + stride_inn * offs_outn[None, :])\n    mm_1_ptrs = mm_1_ptr + (stride_mmm * offs_outm[:, None] + stride_mmn * offs_outn[None, :])\n\n    out_mask = (offs_outm[:, None] < M) & (offs_outn[None, :] < N)\n\n    tl.store(out_ptrs, accumulator, mask=out_mask)\n    tl.store(act_in_ptrs, acc1, mask=out_mask)\n    tl.store(mm_1_ptrs, acc2, mask=out_mask)\n\ndef kernel_ff(x: torch.Tensor, w1: torch.Tensor, w3: torch.Tensor,) -> torch.Tensor:\n    assert x.dtype == torch.float16\n    assert w1.dtype == w3.dtype\n    assert w1.dtype\n    assert w1.shape == w3.shape\n\n    M, K = x.shape\n\n    N = w1.shape[1]\n    assert K == w1.shape[0]\n    assert w1.shape == w3.shape\n    out = torch.empty((M, N), dtype=x.dtype, device=x.device)\n    act_in = torch.empty((M, N), dtype=x.dtype, device=x.device)\n    mm_1 = torch.empty((M, N), dtype=x.dtype, device=x.device)\n\n    grid = lambda META: (triton.cdiv(META[\"M\"], META[\"BLOCK_SIZE_M\"]) * triton.cdiv(META[\"N\"], META[\"BLOCK_SIZE_N\"]),)\n    ff_llama[grid](\n        x, w1, w3, \n        out, act_in, mm_1,\n        M, N, K,\n        *x.stride(),\n        *w1.stride(),\n        *w3.stride(),\n        *out.stride(),\n        *act_in.stride(),\n        *mm_1.stride(),\n    )\n    return out, act_in, mm_1\n\nx = torch.randn([1, 16, 4096], dtype=torch.float16, device=\"cuda\")\nw1_w = torch.randn([11008, 4096], dtype=torch.float16, device=\"cuda\") * 0.2\nw3_w = torch.randn([11008, 4096], dtype=torch.float16, device=\"cuda\") * 0.2\n\nif __name__ == \"__main__\":\n    output_tri, act_in_tri, mm_1_tri = kernel_ff(x=x, w1=w1_w, w3=w3_w)\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "quantize_kv_copy.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/quantize_kv_copy.py_gen_triton_code_136866.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rope_embedding.py": {"call_err_msg": "", "exe_err_msg": "Code is empty", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matrix_vector_multip.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/matrix_vector_multip.py_gen_triton_code_534671.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rms_rbe_matmul.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/rms_rbe_matmul.py_gen_triton_code_448869.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "layer_norm_liger.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/layer_norm_liger.py_gen_triton_code_508425.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "context_attn_llama.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/context_attn_llama.py_gen_triton_code_886853.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "cross_entropy1.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/cross_entropy1.py_gen_triton_code_177582.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "masked_select.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/masked_select.py_gen_triton_code_318627.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "context_attn_mistral.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/context_attn_mistral.py_gen_triton_code_117535.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "attn_fwd_causal.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/attn_fwd_causal.py_gen_triton_code_772187.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "bgmv_shrink_kernel.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/bgmv_shrink_kernel.py_gen_triton_code_881896.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "token_softmax_llama.py": {"call_err_msg": "", "exe_err_msg": "python3: can't open file '/wekafs/zihao/exp/agent/jianghui_triton/TB-eval/tb_eval/../outputs/1107_oss_120b_medium_tmp/tmp/gen/token_softmax_llama.py_gen_triton_code_2914.py': [Errno 2] No such file or directory", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "triton_matmul.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "bmm_optimized.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "triton_attention.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rotary_emb_nopad.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "reversed_cumsum_scalar.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "flash_attn.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "int8_quantization.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "attn_fwd_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rotary_transform_ops.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "softmax_triton1.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "layer_norm_welfold.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "mul_exponent_compensator.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "dropout_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "quantize_copy_kv.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n@triton.jit\ndef sample_from_prob_kernel(x_ptr,  # *Pointer* to first input vector.\n                            output_ptr,  # *Pointer* to output vector.\n                            n_elements,  # Size of the vector.\n                            BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process.\n                            ):\n  \n  # There are multiple 'programs' processing different data. We identify which program\n  # we are here:\n  pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.\n\n  # This program will process inputs that are offset from the initial data.\n  # For instance, if you had a vector of length 256 and block_size of 64, the programs\n  # would each access the elements [0:64, 64:128, 128:192, 192:256].\n  # Note that offsets is a list of pointers:\n  block_start = pid * BLOCK_SIZE\n  offsets = block_start + tl.arange(0, BLOCK_SIZE)\n  # Create a mask to guard memory operations against out-of-bounds accesses.\n  mask = offsets < n_elements\n\n  # Load x and y from DRAM, masking out any extra elements in case the input is not a\n  # multiple of the block size.\n  x = tl.load(x_ptr + offsets, mask=mask)\n  output = x*2\n\n  # Write x + y back to DRAM.\n  tl.store(output_ptr + offsets, output, mask=mask)\n\ndef sample_from_prob(x: torch.Tensor):\n  # We need to preallocate the output.\n  output = torch.empty_like(x)\n  n_elements = output.numel()\n  # The SPMD launch grid denotes the number of kernel instances that run in parallel.\n  # It is analogous to CUDA launch grids. It can be either Tuple[int], or Callable(metaparameters) -> Tuple[int].\n  # In this case, we use a 1D grid where the size is the number of blocks:\n  grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n  # NOTE:\n  #  - Each torch.tensor object is implicitly converted into a pointer to its first element.\n  #  - `triton.jit`'ed functions can be indexed with a launch grid to obtain a callable GPU kernel.\n  #  - Don't forget to pass meta-parameters as keywords arguments.\n  sample_from_prob_kernel[grid](x, output, n_elements, BLOCK_SIZE=1024)\n  # We return a handle to z but, since `torch.cuda.synchronize()` hasn't been called, the kernel is still\n  # running asynchronously at this point.\n  return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "token_attn_mistral.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "var_len_copy.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "attention_score.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "batched_vecmat_mult.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "cross_entropy2.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "context_attn_fwd.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "masked_add_cuda.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rope_transform.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n@triton.jit\ndef _triton_rope(\n    q_ptr,\n    q_row_stride,\n    k_ptr,\n    k_row_stride,\n    cos,\n    cos_row_stride,\n    sin,\n    sin_row_stride,\n    sl,\n    bs: tl.constexpr,\n    n_qh: tl.constexpr,\n    n_kh: tl.constexpr,\n    hd: tl.constexpr,\n    pad_n_qh: tl.constexpr,\n    pad_n_kh: tl.constexpr,\n    pad_hd: tl.constexpr,\n    BLOCK_SIZE: tl.constexpr,\n    BACKWARD_PASS: tl.constexpr = False,\n):\n    # q size: (bsz, seq_len, num_q_heads, head_dim)\n    # q stride: (seq_len * num_q_heads * head_dim, num_q_heads * head_dim, head_dim, 1)\n    # k size: (bsz, seq_len, num_kv_heads, head_dim)\n    # k stride: (seq_len * num_kv_heads * head_dim, num_kv_heads * head_dim, head_dim, 1)\n\n    # cos size: (1, seq_len, head_dim)\n    # stride: (seq_len * head_dim, head_dim, 1)\n    pid = tl.program_id(0)\n\n    # locate start address\n    q_ptr = q_ptr + pid * q_row_stride\n    k_ptr = k_ptr + pid * k_row_stride\n\n    # ####################################################################\n    # get the cos(m\u03b8_{i...d/2}) and sin(m\u03b8_{i...d/2}) for token position\n    # m of this program instance\n    # ####################################################################\n\n    # 1. program instances are laid out in a 1D vector of size bsz * seq_len, which\n    # effectively represents a 2D grid of size [bsz, seq_len] with seq_len dimension\n    # being the fastest changing dimension. Thus we can simply do pid // sl to get the batch index\n    # and pid % sl to get the sequence index.\n    # 2. We only need the left half of cos and sin matrix because the right half is just\n    # a clone of the left half.\n    cos_row_idx = pid % (sl)\n    cos = cos + cos_row_idx * cos_row_stride\n    sin = sin + cos_row_idx * sin_row_stride\n    cos_offsets = tl.arange(0, pad_hd // 2)\n    cos_mask = cos_offsets < hd // 2\n    cos_row = tl.load(cos + cos_offsets, mask=cos_mask, other=0)\n    sin_row = tl.load(sin + cos_offsets, mask=cos_mask, other=0)\n\n    # ####################################################################\n    # Load the left and right half of q and k for the current\n    # program instance (i.e. for the current token) separately\n    # ####################################################################\n    # left half of the head\n    first_half_q_offsets = (\n        tl.arange(0, pad_n_qh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :]\n    )\n    first_half_k_offsets = (\n        tl.arange(0, pad_n_kh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :]\n    )\n    first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & (\n        tl.arange(0, pad_hd // 2)[None, :] < hd // 2\n    )\n    first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & (\n        tl.arange(0, pad_hd // 2)[None, :] < hd // 2\n    )\n    q_tile_1 = tl.load(q_ptr + first_half_q_offsets, mask=first_q_mask, other=0).to(\n        sin_row.dtype\n    )\n    k_tile_1 = tl.load(k_ptr + first_half_k_offsets, mask=first_k_mask, other=0).to(\n        sin_row.dtype\n    )\n\n    # right half of the head\n    second_half_q_offsets = first_half_q_offsets + (hd // 2)\n    second_half_k_offsets = first_half_k_offsets + (hd // 2)\n    second_q_mask = first_q_mask\n    second_k_mask = first_k_mask\n    q_tile_2 = tl.load(q_ptr + second_half_q_offsets, mask=second_q_mask, other=0).to(\n        sin_row.dtype\n    )\n    k_tile_2 = tl.load(k_ptr + second_half_k_offsets, mask=second_k_mask, other=0).to(\n        sin_row.dtype\n    )\n\n    if not BACKWARD_PASS:\n        # y = [x1, x2] * [cos, cos] + [-x2, x1] * [sin, sin]\n        new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row\n        tl.store(q_ptr + first_half_q_offsets, new_q_tile_1, mask=first_q_mask)\n        new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row\n        tl.store(q_ptr + second_half_q_offsets, new_q_tile_2, mask=second_q_mask)\n\n        new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row\n        tl.store(k_ptr + first_half_k_offsets, new_k_tile_1, mask=first_k_mask)\n        new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row\n        tl.store(k_ptr + second_half_k_offsets, new_k_tile_2, mask=second_k_mask)\n    else:\n        # with some math, we can get:\n        # dy = [dx1, dx2] * [cos, cos] + [-dx2, dx1] * [-sin, -sin]\n        new_q_tile_1 = q_tile_1 * cos_row + q_tile_2 * sin_row\n        tl.store(q_ptr + first_half_q_offsets, new_q_tile_1, mask=first_q_mask)\n        new_q_tile_2 = q_tile_2 * cos_row - q_tile_1 * sin_row\n        tl.store(q_ptr + second_half_q_offsets, new_q_tile_2, mask=second_q_mask)\n\n        new_k_tile_1 = k_tile_1 * cos_row + k_tile_2 * sin_row\n        tl.store(k_ptr + first_half_k_offsets, new_k_tile_1, mask=first_k_mask)\n        new_k_tile_2 = k_tile_2 * cos_row - k_tile_1 * sin_row\n        tl.store(k_ptr + second_half_k_offsets, new_k_tile_2, mask=second_k_mask)\n\n\ndef rope_forward(q, k, cos, sin):\n    # transpose it back to the physical shape because Triton looks at the physical storage\n    # note: q and k are incontiguous before the transformation and will become contiguous after transpose\n    q = q.transpose(1, 2)\n    k = k.transpose(1, 2)\n\n    batch_size, seq_len, n_q_head, head_dim = q.shape\n    n_kv_head = k.shape[2]\n    pad_hd = triton.next_power_of_2(head_dim)\n    pad_n_q_head = triton.next_power_of_2(n_q_head)\n    pad_n_kv_head = triton.next_power_of_2(n_kv_head)\n    BLOCK_SIZE = max(pad_n_q_head, pad_n_kv_head)\n\n    n_row = batch_size * seq_len\n\n    # ensure tensors passed into the kernel are contiguous. It will be no-op if they are already contiguous\n    q = q.contiguous()\n    k = k.contiguous()\n    cos = cos.contiguous()\n    sin = sin.contiguous()\n\n    _triton_rope[(n_row,)](\n        q,\n        q.stride(1),\n        k,\n        k.stride(1),\n        cos,\n        cos.stride(-2),\n        sin,\n        sin.stride(-2),\n        seq_len,\n        batch_size,\n        n_q_head,\n        n_kv_head,\n        head_dim,\n        pad_n_q_head,\n        pad_n_kv_head,\n        pad_hd,\n        BLOCK_SIZE=BLOCK_SIZE,\n        BACKWARD_PASS=False,\n    )\n    return q.transpose(1, 2), k.transpose(1, 2), cos, sin\n\n\ndef rope_backward(dq, dk, cos, sin):\n    dq = dq.transpose(1, 2)\n    dk = dk.transpose(1, 2)\n\n    batch_size, seq_len, n_q_head, head_dim = dq.shape\n    n_kv_head = dk.shape[2]\n    pad_hd = triton.next_power_of_2(head_dim)\n    pad_n_q_head = triton.next_power_of_2(n_q_head)\n    pad_n_kv_head = triton.next_power_of_2(n_kv_head)\n    BLOCK_SIZE = max(pad_n_q_head, pad_n_kv_head)\n\n    n_row = batch_size * seq_len\n\n    # ensure dq and dk are contiguous\n    dq = dq.contiguous()\n    dk = dk.contiguous()\n\n    # backward is similar to forward except swapping few ops\n    _triton_rope[(n_row,)](\n        dq,\n        dq.stride(1),\n        dk,\n        dk.stride(1),\n        cos,\n        cos.stride(-2),\n        sin,\n        sin.stride(-2),\n        seq_len,\n        batch_size,\n        n_q_head,\n        n_kv_head,\n        head_dim,\n        pad_n_q_head,\n        pad_n_kv_head,\n        pad_hd,\n        BLOCK_SIZE=BLOCK_SIZE,\n        BACKWARD_PASS=True,\n    )\n    return dq.transpose(1, 2), dk.transpose(1, 2)\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "bmm_chunk_bwd.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "kldiv_ops.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "softmax_triton2.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "pow_scalar_tensor.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "l2_norm_triton2.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fast_layernorm.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "mean_reduction.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "kv_cache_filling.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "relu_triton_kernel.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "max_reduction.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "streamk_matmul.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunk_gate_recurrence.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunk_gated_attention.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "swiglu_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "vector_addition_custom.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import triton\nimport torch\n\n# Triton kernel for element-wise addition\n@triton.jit\ndef add_kernel(X, Y, Z, N, BLOCK_SIZE: tl.constexpr):\n    pid = tl.program_id(0)\n    block_start = pid * BLOCK_SIZE\n    offsets = block_start + tl.arange(0, BLOCK_SIZE)\n    mask = offsets < N\n    x = tl.load(X + offsets, mask=mask)\n    y = tl.load(Y + offsets, mask=mask)\n    z = x + y\n    tl.store(Z + offsets, z, mask=mask)\n\n# Function to call the Triton kernel\ndef add_tensors(x: torch.Tensor, y: torch.Tensor):\n    assert x.is_cuda and y.is_cuda\n    assert x.numel() == y.numel()\n    z = torch.empty_like(x)\n    BLOCK_SIZE = 1024\n    grid = lambda meta: (triton.cdiv(x.numel(), meta['BLOCK_SIZE']),)\n    add_kernel[grid](x, y, z, x.numel(), BLOCK_SIZE=BLOCK_SIZE)\n    return z\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "triton_conv2d_fwd.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "isfinite_kernel.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fp4_to_bf16.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rope_backward_transform.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "sgmv_expand_slice.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "embedding_triton_kernel.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "uniform_sampling.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "bgmv_expand_slice.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fused_recurrent_delta.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "vector_addition.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rms_norm_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "multinomial_sampling.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\nfrom flag_gems.utils.random_utils import philox_cuda_seed_offset, uniform\nfrom flag_gems.ops import normed_cumsum\n\n@triton.jit(do_not_specialize=[\"K\", \"N\", \"philox_seed\", \"philox_offset\"])\ndef multinomial_with_replacement(\n    cdf_ptr, out_ptr, K, N, philox_seed, philox_offset, NBLOCK: tl.constexpr\n):\n    # The computation is arranged in a 2d grid of blocks, each producing\n    # a batch of samples for a particular distribution.\n    y_off = tl.program_id(1) * N\n    n = tl.program_id(0) * NBLOCK + tl.arange(0, NBLOCK)\n    rv, _, _, _ = uniform(philox_seed, philox_offset, y_off + n)\n\n    # Do a binary search for each random number on the cumulative probabilities.\n    rv += 0.0001\n    rv = tl.where(rv > 0.9999, 0.9999, rv)\n\n    cdf_ptr += tl.program_id(1) * K\n    start = tl.zeros((NBLOCK,), dtype=tl.int32)\n    end = tl.zeros((NBLOCK,), dtype=tl.int32) + K - 1\n    steps = tl.math.log2(K.to(tl.float32)).to(tl.int32) + 1\n    for _ in range(steps):\n        mid = start + (end - start) // 2\n        x = tl.load(cdf_ptr + mid, mask=n < N)\n        start = tl.where(x < rv, mid + 1, start)\n        end = tl.where(x < rv, end, mid)\n\n    # Returns the last index in case of an overflow\n    start = tl.where(start >= K, K - 1, start)\n\n    tl.store(out_ptr + y_off + n, start, mask=n < N)\n\ndef multinomial(prob, n_samples, with_replacement=False, *, gen=None):\n    assert prob.dtype in (torch.float16, torch.float32, torch.bfloat16, torch.float64)\n    assert 0 < prob.dim() <= 2, \"prob_dist must be 1 or 2 dim\"\n    n_categories = prob.size(-1)\n    assert n_categories <= (1 << 24), \"number of categories cannot exceed 2^24\"\n    assert (\n        with_replacement or n_samples <= n_categories\n    ), \"cannot sample n_samples > prob.size(-1) samples without replacement.\"\n\n    # Sampling without replacement\n    if (not with_replacement) or n_samples == 1:\n        q = torch.empty_like(prob).exponential_(1.0)\n        s = torch.div(prob, q, out=q)\n        if n_samples == 1:\n            return torch.argmax(s, dim=-1, keepdim=True).to(torch.int64)\n        else:\n            vals, indices = torch.topk(s, n_samples, dim=-1)\n            return indices.to(torch.int64)\n\n    cum_prob = normed_cumsum(prob, dim=-1)\n\n    if cum_prob.dim() == 1:\n        n_dist = 1\n        out = torch.empty((n_samples,), device=prob.device, dtype=torch.int64)\n    else:\n        n_dist = cum_prob.size(0)\n        out = torch.empty((n_dist, n_samples), device=prob.device, dtype=torch.int64)\n    \n    increment = n_dist * n_samples\n    philox_seed, philox_offset = philox_cuda_seed_offset(increment)\n    grid = lambda META: (triton.cdiv(n_samples, META[\"NBLOCK\"]), n_dist)\n    multinomial_with_replacement[grid](\n        cum_prob, out, n_categories, n_samples, philox_seed, philox_offset\n    )\n    return out\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "mixed_sparse_attention.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunk_retention_ops.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "square_matrix.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunk_gla_fwd.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "layer_norm_ops.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "kv_cache_copy.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunk_gla_simple.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "int8_matmul_quantization.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "seeded_dropout.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "int8_dequant_matmul.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "quant_transpose_kernel.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "token_attn_reduceV.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\ndef test_triton_heuristic(device):\n    N = 1023\n    src = torch.empty(N, device=device)\n    dst = torch.zeros(N, device=device)\n\n    do_bench = lambda kernel, quantiles: triton.testing.do_bench(kernel, quantiles=quantiles, warmup=1, rep=1)\n\n    # Triton kernel with @triton.jit decorator\n    @triton.autotune(configs=[triton.Config(kwargs={'BLOCK_SIZE': 32})], key=['N'], do_bench=do_bench)\n    @triton.heuristics({'EVEN_N': lambda nargs: nargs['N'] % 2 == 0})  # test kwargs\n    @triton.heuristics({'EVEN_src': lambda nargs: nargs['src'].data_ptr() % 2 == 0})  # test args\n    @triton.jit\n    def _kernel(dst, src, N, BLOCK_SIZE: tl.constexpr, EVEN_N: tl.constexpr, EVEN_src: tl.constexpr):\n        # Kernel logic\n        tl.store(dst, EVEN_N)\n        tl.store(dst + 1, EVEN_src)\n\n    grid = lambda META: (triton.cdiv(N, META['BLOCK_SIZE']), )\n    _kernel[grid](dst, src, N=N)\n    assert dst[0].item() == 0.0\n    assert dst[1].item() == 1.0\n    assert _kernel.base_fn.__name__ == \"_kernel\"\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "layer_norm_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matrix_reduction.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "layernorm_fwd_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "attention_kernel.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "flash_decode2_llama.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "index_select_bwd.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "kldiv_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matmul_dequantize.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rowwise_quantization_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "softmax_reducev.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "int4_matmul.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "add_value.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "f8_conversion_utils.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matmul_kernel.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fused_recurrent_retention.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rmsnorm_fused_llama.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "destindex_copy_kv1.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rotary_emb.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "flash_decode2_phi.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matmul_triton_autotune.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "cache_transform_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "parallel_retention_attention.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fast_rms_layernorm.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "attention_fwd_triton2.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rbe_triton_transform.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "iv_dependent_matmul.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "swiglu_backward.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import triton\nimport triton.language as tl\nimport torch\n\n# Define Triton kernel\n@triton.jit\ndef my_kernel(X, Y, Z, N):\n    idx = tl.arange(0, N)\n    Z[idx] = X[idx] + Y[idx]\n\n# Function to call the kernel\ndef call_my_kernel(X, Y, Z, N):\n    # Obtain necessary pointers\n    x_ptr = X.data_ptr()\n    y_ptr = Y.data_ptr()\n    z_ptr = Z.data_ptr()\n\n    # Launch the Triton kernel\n    grid = lambda opt: (triton.cdiv(N, opt.d('BLOCK')),)\n    my_kernel[grid](X, Y, Z, N)\n\n# Dummy inputs for the kernel call\nX = torch.randn(1024, device='cuda')\nY = torch.randn(1024, device='cuda')\nZ = torch.empty_like(X)\nN = X.numel()\n\n# Invoke the kernel\ncall_my_kernel(X, Y, Z, N)\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matrix_transpose.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n@triton.jit\ndef _fwd_kernel_destindex_copy_quantize_int4_kv(\n    K,\n    Dest_loc,\n    Out,\n    Out_scale,\n    stride_k_bs,\n    stride_k_h,\n    stride_k_g,\n    stride_k_d,\n    stride_o_bs,\n    stride_o_h,\n    stride_o_g,\n    stride_o_d,\n    stride_os_bs,\n    stride_os_h,\n    stride_os_g,\n    group_size,\n    BLOCK_GROUP_NUM: tl.constexpr,\n    BLOCK_GROUP_DIM: tl.constexpr,\n):\n    cur_index = tl.program_id(0)\n    cur_head = tl.program_id(1)\n\n    offs_g = tl.arange(0, BLOCK_GROUP_NUM)\n    offs_d = tl.arange(0, BLOCK_GROUP_DIM // 2)\n\n    dest_index = tl.load(Dest_loc + cur_index)\n\n    src_data_0 = tl.load(\n        K + cur_index * stride_k_bs + cur_head * stride_k_h + offs_g[:, None] * stride_k_g + offs_d[None, :] * 2,\n        mask=offs_g[:, None] < group_size,\n        other=0.0,\n    )\n    src_data_1 = tl.load(\n        K + cur_index * stride_k_bs + cur_head * stride_k_h + offs_g[:, None] * stride_k_g + offs_d[None, :] * 2 + 1,\n        mask=offs_g[:, None] < group_size,\n        other=0.0,\n    )\n\n    abs_data_0 = tl.abs(src_data_0)\n    abs_data_1 = tl.abs(src_data_1)\n\n    data_scale = (tl.maximum(tl.max(abs_data_0, axis=1), tl.max(abs_data_1, axis=1)) / 7.0).to(Out_scale.dtype.element_ty)\n    q_src_data_0 = (src_data_0 / data_scale[:, None]).to(tl.int8)\n    q_src_data_0 = tl.where(q_src_data_0 > 7, 7, q_src_data_0)\n    q_src_data_0 = tl.where(q_src_data_0 < -7, -7, q_src_data_0)\n\n    q_src_data_1 = (src_data_1 / data_scale[:, None]).to(tl.int8)\n    q_src_data_1 = tl.where(q_src_data_1 > 7, 7, q_src_data_1)\n    q_src_data_1 = tl.where(q_src_data_1 < -7, -7, q_src_data_1)\n\n    low_4 = ((q_src_data_0 & 0x80) >> 4) | (q_src_data_0 & 0xF)\n    high_4 = (((q_src_data_1 & 0x80) >> 4) | (q_src_data_1 & 0xF)) << 4\n\n    out_data = low_4 | high_4\n\n    o_ptrs = Out + dest_index * stride_o_bs + cur_head * stride_o_h + offs_g[:, None] * stride_o_g + offs_d[None, :]\n    os_ptrs = Out_scale + dest_index * stride_os_bs + cur_head * stride_os_h + offs_g\n    tl.store(o_ptrs, out_data, mask=offs_g[:, None] < group_size)\n    tl.store(os_ptrs, data_scale, mask=offs_g < group_size)\n    return\n\n@torch.no_grad()\ndef destindex_copy_int4kv(K, DestLoc, Out, Out_scale):\n    head_dim = K.shape[2]\n    quant_group_dim = 8\n\n    assert head_dim % quant_group_dim == 0, \"error head dim, can not been supported to copy quant kv\"\n\n    group_size = head_dim // quant_group_dim\n    group_dim = quant_group_dim\n\n    K = K.view((K.shape[0], K.shape[1], group_size, group_dim))\n    Out = Out.view(\n        Out.shape[0], Out.shape[1], group_size, group_dim // 2\n    )\n\n    # _fwd_kernel_destindex_copy_quantize_int4_kv[grid](\n    #     K,\n    #     DestLoc,\n    #     Out,\n    #     Out_scale,\n    #     K.stride(0),\n    #     K.stride(1),\n    #     K.stride(2),\n    #     K.stride(3),\n    #     Out.stride(0),\n    #     Out.stride(1),\n    #     Out.stride(2),\n    #     Out.stride(3),\n    #     Out_scale.stride(0),\n    #     Out_scale.stride(1),\n    #     Out_scale.stride(2),\n    #     group_size,\n    #     BLOCK_GROUP_NUM=triton.next_power_of_2(group_size),\n    #     BLOCK_GROUP_DIM=group_dim,\n    #     num_warps=num_warps,\n    #     num_stages=1,\n    # )\n    return\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "softmax_optimize.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "lora_expand_gemv.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "attention_fwd_triton3.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "destindex_copy.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fused_layernorm_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "spinning_lock_reduction.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rmsnorm_implementation.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunk_bwd_dqkg.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fifth_order_sph_harmonics.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fused_rwkv6_kernel.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "triton_softmax.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunk_cumsum_kernel.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "relu_strided_buffer.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fast_rope_embedding.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "reversed_cumsum.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import triton\nimport triton.language as tl\nimport torch\n\n# Triton kernel for streaming logsumexp\n@triton.jit\ndef _logsumexp(X, OUT, xm_stride, xn_stride, out_stride, N, BLOCK_N: tl.constexpr):\n    rm = tl.program_id(0)\n    alpha = tl.zeros((1,), tl.float32) + -float('inf')\n    res = tl.zeros((1,), tl.float32)\n    for bn in range(0, N, BLOCK_N):\n        rn = bn + tl.arange(0, BLOCK_N)\n        Xmn = X + rm * xm_stride + rn * xn_stride\n        x = tl.load(Xmn, mask=rn < N, other=-float('inf'))\n        c = tl.max(x, axis=0)\n        # correct the current sum and update the max\n        res = tl.where(c > alpha, res * tl.exp(alpha - c), res)\n        alpha = tl.where(c > alpha, c, alpha)\n        res += tl.sum(tl.exp(x - alpha), axis=0)\n    out = tl.log(res) + alpha\n    rm = tl.program_id(0) + tl.arange(0, 1)\n    OUT = OUT + rm * out_stride\n    tl.store(OUT, out)\n\n# Function to call the Triton kernel\ndef logsumexp(input):\n    assert input.is_cuda\n    *dims, N = input.shape\n    input = input.view(-1, N)\n    out = input.new_empty(*dims).view(-1)\n    M = input.shape[0]\n    _logsumexp[(M,)](input, out, input.stride(0), input.stride(1), out.stride(0), N,\n                     BLOCK_N=4096, num_warps=4)\n    return out.view(*dims)\n\n# Softmax function using the logsumexp Triton kernel\ndef softmax_(x):\n    if not x.is_cuda:\n        return torch.softmax(x, dim=-1, out=x)\n    c = logsumexp(x)\n    return x.sub_(c[..., None]).exp_()\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "logsumexp_fwd.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "parallel_attention.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matmul_dequant_int4.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "quantize_kv_transform.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rotary_transform.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "apply_penalty.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunk_delta_fwd.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matmul_leakyrelu.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "attention_llama.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "destindex_copy_kv2.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "context_attn_nopad.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "sin_kernel.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "cross_entropy_ops.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "add_example.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "block_sparse_attn.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "int_scaled_matmul.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "decay_cumsum.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matmul_persistent_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "geglu_tanh_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "bmm_chunk_fwd.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunked_cumsum_fwd.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "dequantize_matmul.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fused_activation.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matmul_tma.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "triton_mul2.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "adam_update_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "triton_argmax.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "layer_norm_fwd.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunk_retention.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "sin_computation.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fused_rotary_embedding.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunk_linear_attn.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "ksoftmax_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rmsnorm_fused.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rmsnorm_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "nested_loops_processing.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matmul_dequantize_int4.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fp4_to_bf16_conversion.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matmul_triton1.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matmul_leakyrelu_fp8.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "triton_linear_activation.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "softmax_triton3.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "attention_fwd_triton1.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "attention_forward_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "matmul_triton2.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "kcache_copy_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "cosine_compute.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fast_ce_loss.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "log_softmax.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "index_select_cat.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n# Define the Triton kernel for flattening a 2D matrix to a 1D array\n@triton.jit\ndef flatten_kernel(\n    input_matrix,   # Pointer to the input matrix in global memory\n    output_array,   # Pointer to the output array in global memory\n    rows,           # Number of rows in the input matrix\n    cols,           # Number of columns in the input matrix\n    BLOCK_SIZE: tl.constexpr,\n):\n    # Compute row and column index for this thread\n    row = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)\n    col = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)\n\n    # Ensure we do not go out of bounds\n    mask = (row < rows) & (col < cols)\n\n    # Compute the linear index in the flattened output array\n    idx = row * cols + col\n\n    # Load from the input matrix and store in the output array\n    val = tl.load(input_matrix + row[:, None] * cols + col[None, :], mask=mask)\n    tl.store(output_array + idx, val, mask=mask)\n\n# Setup the PyTorch tensors\nrows = 64\ncols = 64\ninput_matrix = torch.randn(rows, cols, dtype=torch.float32, device='cuda')\noutput_array = torch.empty(rows * cols, dtype=torch.float32, device='cuda')\n\n# Define grid and block sizes for the kernel\nBLOCK_SIZE = 16  # Define according to your GPU's capability and the size of the matrix\ngrid = (rows // BLOCK_SIZE, cols // BLOCK_SIZE)\n\n# Launch the kernel\nflatten_kernel[grid,](\n    input_matrix, output_array,\n    rows, cols,\n    BLOCK_SIZE=BLOCK_SIZE\n)\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "token_attn_llama2.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _fwd_kernel_token_att1(\n    Q,\n    K,\n    sm_scale,\n    Req_to_tokens,\n    B_req_idx,\n    B_Start_Loc,\n    B_Seqlen,\n    B_Att_Start_Loc,\n    B_Att_Seqlen,\n    Att_Out,\n    stride_req_to_tokens_b,\n    stride_req_to_tokens_s,\n    stride_qbs,\n    stride_qh,\n    stride_qd,\n    stride_kbs,\n    stride_kh,\n    stride_kd,\n    att_stride_h,\n    att_stride_bs,\n    kv_group_num,\n    sliding_window,\n    BLOCK_DMODEL: tl.constexpr,\n    BLOCK_N: tl.constexpr,\n):\n    cur_batch = tl.program_id(0)\n    cur_head = tl.program_id(1)\n    start_n = tl.program_id(2)\n\n    cur_kv_head = cur_head // kv_group_num\n\n    offs_d = tl.arange(0, BLOCK_DMODEL)  # [D]\n    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)\n    cur_batch_in_all_start_index = tl.load(B_Att_Start_Loc + cur_batch)  # use window index\n    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)\n    cur_att_seq_len = tl.load(B_Att_Seqlen + cur_batch)\n\n    # use new start index of k value\n    cur_batch_start_index = tl.maximum(cur_batch_seq_len - sliding_window, 0)\n    cur_batch_end_index = cur_batch_seq_len\n\n    off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d * stride_qd  # [D]\n\n    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)  # [32]\n\n    # use new value to decide block mask\n    block_stard_index = start_n * BLOCK_N\n    block_mask = tl.where(block_stard_index < cur_att_seq_len, 1, 0)  # a number\n\n    for start_mark in range(0, block_mask, 1):\n        q = tl.load(Q + off_q + start_mark)  # [SYM] why here add start_mark\n        offs_n_new = cur_batch_start_index + offs_n  # the latest window of token\n        k_loc = tl.load(\n            Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * offs_n_new,\n            mask=offs_n_new < cur_batch_end_index,\n            other=0,\n        )\n        off_k = (\n            k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :] * stride_kd\n        )  # [32, D], find token index\n        k = tl.load(K + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0)\n        att_value = tl.sum(q[None, :] * k, 1)  # [1, D] * [32, D] = [32, D] -> [32]\n        att_value = att_value.to(tl.float32)\n        att_value *= sm_scale\n        off_o = cur_head * att_stride_h + (cur_batch_in_all_start_index + offs_n) * att_stride_bs\n        tl.store(Att_Out + off_o, att_value, mask=offs_n_new < cur_batch_end_index)\n    return\n\n\n@torch.no_grad()\ndef token_att_fwd(\n    q, k, att_out, Req_to_tokens, B_req_idx, B_Start_Loc, B_Seqlen, B_Att_Start_Loc, B_Att_Seqlen, sliding_window\n):\n    BLOCK = 32\n    # shape constraints\n    Lq, Lk = q.shape[-1], k.shape[-1]\n    assert Lq == Lk\n    assert Lk in {16, 32, 64, 128}\n    sm_scale = 1.0 / (Lk ** 0.5)\n\n    batch, head_num = B_req_idx.shape[0], q.shape[1]\n\n    grid = (batch, head_num, triton.cdiv(sliding_window, BLOCK))\n    kv_group_num = q.shape[1] // k.shape[1]\n\n    if kv_group_num == 1:\n        num_warps = 4\n    else:\n        num_warps = 2\n\n    _fwd_kernel_token_att1[grid](\n        q,\n        k,\n        sm_scale,\n        Req_to_tokens,\n        B_req_idx,\n        B_Start_Loc,\n        B_Seqlen,\n        B_Att_Start_Loc,\n        B_Att_Seqlen,\n        att_out,\n        Req_to_tokens.stride(0),\n        Req_to_tokens.stride(1),\n        q.stride(0),\n        q.stride(1),\n        q.stride(2),\n        k.stride(0),\n        k.stride(1),\n        k.stride(2),\n        att_out.stride(0),\n        att_out.stride(1),\n        kv_group_num=kv_group_num,\n        sliding_window=sliding_window,\n        BLOCK_DMODEL=Lk,\n        BLOCK_N=BLOCK,\n        num_warps=num_warps,\n        num_stages=1,\n    )\n    return\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "l2_norm_bwd.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "token_softmax_bloom.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "chunk_cumsum_vector.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "diag_ssm_triton.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "quantize_global.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import math\nimport torch\nimport triton\nimport triton.language as tl\n\n# This kernel does fused columnwise quantization and transpose.\n@triton.jit\ndef _quantize_columnwise_and_transpose(\n    x_ptr,\n    output_ptr,\n    output_maxs,\n    n_elements,\n    M: tl.constexpr,\n    N: tl.constexpr,\n    BLOCK_SIZE: tl.constexpr,\n    P2: tl.constexpr,\n):\n    pid = tl.program_id(axis=0)\n    block_start = pid\n    p2_arange = tl.arange(0, P2)\n    p2_arange_mask = p2_arange < M\n    arange = p2_arange * N\n    offsets = block_start + arange\n    x = tl.load(x_ptr + offsets, mask=p2_arange_mask)\n    abs_x = tl.abs(x)\n    max_val = tl.max(tl.where(p2_arange_mask, abs_x, 0), axis=0)\n    output = tl.libdevice.llrint(127.0 * (x / max_val))\n\n    new_start = pid * M\n    new_offsets = new_start + p2_arange\n    tl.store(output_ptr + new_offsets, output, mask=p2_arange_mask)\n    tl.store(output_maxs + pid, max_val)\n\ndef quantize_columnwise_and_transpose(x: torch.Tensor):\n    M, N = x.shape\n    output = torch.empty(N, M, device=x.device, dtype=torch.int8)\n    output_maxs = torch.empty(x.shape[1], device=x.device, dtype=torch.float16)\n\n    P2 = int(2 ** (math.ceil(math.log2(M))))\n\n    assert x.is_cuda and output.is_cuda\n    n_elements = output.numel()\n    grid = lambda meta: (triton.cdiv(n_elements, meta[\"BLOCK_SIZE\"]),)\n    _quantize_columnwise_and_transpose[grid](x, output, output_maxs, n_elements, M, N, BLOCK_SIZE=M, P2=P2)\n    return output, output_maxs\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "l2_norm_triton1.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "rms_matmul_rbe.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "fused_recurrent_hgrn.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef _flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, kv_group_num,\n    x, sm_scale, stride_qt, stride_qh, stride_qd, stride_kcb, stride_kch, stride_kcsplit_x, stride_kcs,\n    stride_kcd, stride_vcb, stride_vch, stride_vcs, stride_vcd, stride_bts, stride_btb, stride_mid_ot,\n    stride_mid_oh, stride_mid_ob, stride_mid_od, stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_block = tl.arange(0, BLOCK_SIZE)\n\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    cur_kv_head_idx = cur_head_idx // kv_group_num\n    offset_kvcache = cur_block_id * stride_kcb + cur_kv_head_idx * stride_kch\n    offsets_k = (\n        offset_kvcache\n        + (offsets_dmodel[None, :] // x) * stride_kcsplit_x\n        + (offsets_dmodel[None, :] % x) * stride_kcd\n        + offsets_block[:, None] * stride_kcs\n    )\n    k_cur_block = tl.load(KCache + offsets_k)\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_vcs, stride_vcd),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij += tl.where(block_start_kv * BLOCK_KV + offsets_block < cur_kv_seq_len, 0, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _alibi_flash_decoding_fwd_kernel(\n    Q, KCache, VCache, block_tables, mid_o, mid_o_lse, kv_seq_len, q_len, batch_size, alibi_slopes,\n    stride_qt, stride_qh, stride_qd, stride_cacheb, stride_cacheh, stride_cachebs, stride_cached,\n    stride_bts, stride_btb, stride_mid_ot, stride_mid_oh, stride_mid_ob, stride_mid_od,\n    stride_mid_o_lset, stride_mid_o_lseh, stride_mid_o_lseb, sm_scale, KV_GROUPS: tl.constexpr,\n    BLOCK_KV: tl.constexpr, BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_head_idx = tl.program_id(1)\n    block_start_kv = tl.program_id(2)\n\n    tl.static_assert(BLOCK_KV == BLOCK_SIZE)\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    if block_start_kv * BLOCK_KV >= cur_kv_seq_len:\n        return\n\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n    offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd\n    q = tl.load(Q + offsets_q)\n    block_table_ptr = block_tables + cur_seq_idx * stride_bts\n    cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb)\n    cur_occupied_size = tl.where(\n        (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE\n    )\n    tl.device_assert(cur_occupied_size >= 0)\n\n    cur_kv_head_idx = cur_head_idx // KV_GROUPS\n    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh\n    K_block_ptr = tl.make_block_ptr(\n        base=KCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        base=VCache + offset_kvcache,\n        shape=(cur_occupied_size, HEAD_DIM),\n        strides=(stride_cachebs, stride_cached),\n        offsets=(0, 0),\n        block_shape=(BLOCK_SIZE, HEAD_DIM),\n        order=(0, 1),\n    )\n    k_cur_block = tl.load(K_block_ptr)\n    v_cur_block = tl.load(V_block_ptr)\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n    S_ij = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n\n    alibi_slope = tl.load(alibi_slopes + cur_head_idx)\n    position_k_offset = block_start_kv * BLOCK_KV + tl.arange(0, BLOCK_SIZE)\n\n    S_ij += tl.sum(q[None, :] * k_cur_block, 1)\n    S_ij *= sm_scale\n    S_ij -= alibi_slope * (cur_kv_seq_len - 1 - position_k_offset)\n    S_ij = tl.where(cur_kv_seq_len > position_k_offset, S_ij, float(\"-inf\"))\n\n    m = tl.max(S_ij, 0)\n    S_ij -= m\n    p_ij_hat = tl.exp(S_ij)\n    l_i = tl.sum(p_ij_hat, 0)\n    p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)\n    acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)\n    acc = acc / l_i\n\n    offsets_mid_o = (\n        cur_token_idx * stride_mid_ot\n        + cur_head_idx * stride_mid_oh\n        + block_start_kv * stride_mid_ob\n        + offsets_dmodel * stride_mid_od\n    )\n    tl.store(mid_o + offsets_mid_o, acc)\n    offsets_mid_o_lse = (\n        cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb\n    )\n    tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))\n\n\n@triton.jit\ndef _flash_decoding_fwd_reduce_kernel(\n    mid_o, mid_o_lse, O, kv_seq_len, q_len, batch_size, stride_mid_ot, stride_mid_oh,\n    stride_mid_ob, stride_mid_od, stride_o_lset, stride_o_lseh, stride_o_lseb,\n    stride_ot, stride_oh, stride_od, BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr,\n):\n    cur_token_idx = tl.program_id(0)\n    cur_seq_idx = cur_token_idx // q_len\n    if cur_seq_idx >= batch_size:\n        return\n    cur_head_idx = tl.program_id(1)\n\n    cur_token_off = (cur_token_idx % q_len) - q_len + 1\n    cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + cur_token_off\n    offsets_dmodel = tl.arange(0, HEAD_DIM)\n\n    kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV\n    m_i = float(\"-inf\")\n    l_i = 0.0\n    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)\n\n    offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel\n    offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh\n    for block_i in range(0, kv_split_num, 1):\n        mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob)\n        lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb)\n        m_ij = tl.maximum(m_i, lse)\n        scale = tl.exp(m_i - m_ij)\n        acc = acc * scale\n        lse -= m_ij\n        exp_logic = tl.exp(lse)\n        acc += exp_logic * mid_o_block\n        l_i = scale * l_i + exp_logic\n        m_i = m_ij\n\n    acc = acc / l_i\n    offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel\n    tl.store(O + offsets_O, acc.to(O.type.element_ty))\n    return\n\n\ndef flash_decoding_attention(\n    q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, kv_seq_len: torch.Tensor,\n    block_tables: torch.Tensor, block_size: int, max_seq_len_in_batch: int = None, output: torch.Tensor = None,\n    mid_output: torch.Tensor = None, mid_output_lse: torch.Tensor = None, alibi_slopes: torch.Tensor = None,\n    sm_scale: int = None, kv_group_num: int = 1, q_len: int = 1, use_new_kcache_layout: bool = False,\n):\n    q = q.squeeze() if q.dim() == 4 else q\n    assert q.dim() == 3, f\"Incompatible q dim: {q.dim()}\"\n    n_tokens, num_heads, head_dim = q.shape\n    assert n_tokens % q_len == 0, \"Invalid q_len\"\n    bsz = n_tokens // q_len\n\n    assert head_dim in {32, 64, 128, 256}\n    assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, (\n        f\"Got incompatible batch size (number of seqs):\\n\"\n        f\"  KV seq lengths bsz {kv_seq_len.size(0)}, Block tables bsz {block_tables.size(0)}, \"\n        f\"batch size {bsz}\"\n    )\n    assert k_cache.size(-2) == v_cache.size(-2) == block_size, (\n        f\"Got incompatible block size on kv caches:\\n\"\n        f\"  assigned block_size {block_size}, k_cache block_size {k_cache.size(-2)}, \"\n        f\"v_cache block_size {v_cache.size(-2)}\"\n    )\n\n    assert block_size in {16, 32, 64, 128}\n    BLOCK_KV = block_size\n\n    sm_scale = 1.0 / (head_dim**0.5) if sm_scale is None else sm_scale\n    max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch\n    kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV\n\n    if mid_output is None:\n        mid_output = torch.empty(\n            (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device\n        )\n    if mid_output_lse is None:\n        mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device)\n    if output is None:\n        output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device)\n\n    assert (\n        mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num\n    ), \"Incompatible kv split number of intermediate output tensors\"\n    assert (\n        mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens\n    ), f\"Incompatible first dimension of output tensors\"\n\n    grid = lambda META: (\n        triton.next_power_of_2(bsz * q_len),\n        num_heads,\n        triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), META[\"BLOCK_KV\"]),\n    )\n\n    if alibi_slopes is not None:\n        assert (\n            not use_new_kcache_layout\n        ), \"Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready\"\n\n        _alibi_flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            alibi_slopes,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            k_cache.stride(2),\n            k_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            sm_scale,\n            KV_GROUPS=kv_group_num,\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n    else:\n        x = head_dim\n        kcsplit_x_stride, kcs_stride, kcd_stride = 0, k_cache.stride(2), k_cache.stride(3)\n        if use_new_kcache_layout:\n            assert (\n                k_cache.dim() == 5\n                and k_cache.shape[1] == v_cache.shape[1]\n                and k_cache.shape[2] * k_cache.shape[4] == v_cache.shape[3]\n            ), f\"Invalid KCache shape {k_cache.shape} and VCache shape {v_cache.shape}\"\n            x = k_cache.size(-1)\n            kcsplit_x_stride, kcs_stride, kcd_stride = k_cache.stride()[-3:]\n\n        _flash_decoding_fwd_kernel[grid](\n            q,\n            k_cache,\n            v_cache,\n            block_tables,\n            mid_output,\n            mid_output_lse,\n            kv_seq_len,\n            q_len,\n            bsz,\n            kv_group_num,\n            x,\n            sm_scale,\n            q.stride(0),\n            q.stride(1),\n            q.stride(2),\n            k_cache.stride(0),\n            k_cache.stride(1),\n            kcsplit_x_stride,\n            kcs_stride,\n            kcd_stride,\n            v_cache.stride(0),\n            v_cache.stride(1),\n            v_cache.stride(2),\n            v_cache.stride(3),\n            block_tables.stride(0),\n            block_tables.stride(1),\n            mid_output.stride(0),\n            mid_output.stride(1),\n            mid_output.stride(2),\n            mid_output.stride(3),\n            mid_output_lse.stride(0),\n            mid_output_lse.stride(1),\n            mid_output_lse.stride(2),\n            BLOCK_KV=block_size,\n            BLOCK_SIZE=block_size,\n            HEAD_DIM=head_dim,\n        )\n\n    grid = (triton.next_power_of_2(bsz * q_len), num_heads)\n    _flash_decoding_fwd_reduce_kernel[grid](\n        mid_output,\n        mid_output_lse,\n        output,\n        kv_seq_len,\n        q_len,\n        bsz,\n        mid_output.stride(0),\n        mid_output.stride(1),\n        mid_output.stride(2),\n        mid_output.stride(3),\n        mid_output_lse.stride(0),\n        mid_output_lse.stride(1),\n        mid_output_lse.stride(2),\n        output.stride(0),\n        head_dim,\n        1,\n        BLOCK_KV=block_size,\n        HEAD_DIM=head_dim,\n    )\n\n    return output\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}, "softmax_flaggems.py": {"call_err_msg": "None", "exe_err_msg": "None", "reflection": null, "oneshot": "import torch\nimport triton\nimport triton.language as tl\n\ndef heur_tile_k(args):\n    tile_k = 1\n    MAX_TILE_K = 8192\n    NUM_SMS = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count\n    upper_bound = min(args[\"K\"], MAX_TILE_K)\n    while tile_k <= upper_bound:\n        num_blocks = args[\"M\"] * triton.cdiv(args[\"K\"], tile_k)\n        num_waves = num_blocks / NUM_SMS\n        if (num_waves > 1) and (tile_k * 2 <= upper_bound):\n            tile_k *= 2\n        else:\n            break\n    return tile_k\n\ndef heur_tile_n_non_inner(args):\n    return triton.cdiv(8192, args[\"TILE_K\"])\n\ndef heur_one_tile_per_cta(args):\n    return args[\"TILE_N\"] >= args[\"N\"]\n\ndef heur_num_warps_non_inner(args):\n    tile_size = args[\"TILE_N\"] * args[\"TILE_K\"]\n    if tile_size < 2048:\n        return 4\n    elif tile_size < 4096:\n        return 8\n    else:\n        return 16\n\n@triton.heuristics(\n    {\n        \"TILE_K\": heur_tile_k,\n        \"TILE_N\": heur_tile_n_non_inner,\n        \"ONE_TILE_PER_CTA\": heur_one_tile_per_cta,\n        \"num_warps\": heur_num_warps_non_inner,\n    }\n)\n@triton.jit\ndef softmax_kernel_non_inner(\n    output_ptr,\n    input_ptr,\n    M,\n    N,\n    K,\n    TILE_N: tl.constexpr,\n    TILE_K: tl.constexpr,\n    ONE_TILE_PER_CTA: tl.constexpr,\n):\n    # Kernel logic here...\n    pass\n\ndef heur_tile_n_inner(args):\n    if args[\"N\"] <= (32 * 1024):\n        return triton.next_power_of_2(args[\"N\"])\n    else:\n        return 4096\n\ndef heur_num_warps_inner(args):\n    tile_size = args[\"TILE_N\"]\n    if tile_size < 2048:\n        return 4\n    elif tile_size < 4096:\n        return 8\n    else:\n        return 16\n\n@triton.heuristics(\n    {\n        \"TILE_N\": heur_tile_n_inner,\n        \"ONE_TILE_PER_CTA\": heur_one_tile_per_cta,\n        \"num_warps\": heur_num_warps_inner,\n    }\n)\n@triton.jit\ndef softmax_kernel_inner(\n    output_ptr,\n    input_ptr,\n    M,\n    N,\n    TILE_N: tl.constexpr,\n    ONE_TILE_PER_CTA: tl.constexpr,\n):\n    # Kernel logic here...\n    pass\n\nclass Softmax(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, x, dim, dtype):\n        assert dim >= -x.ndim and dim < x.ndim, \"Invalid dim\"\n        dim = dim % x.ndim\n        M = 1\n        N = x.shape[dim]\n        for i in range(dim):\n            M *= x.shape[i]  # pre_dim\n        inp = x.contiguous()\n        if dtype is None:\n            dtype = x.dtype\n        out = torch.empty_like(inp, dtype=dtype)\n        K = inp.numel() // M // N  # post_dim\n\n        with torch.cuda.device(inp.device):\n            if K > 1:\n                grid = lambda meta: (M, triton.cdiv(K, meta[\"TILE_K\"]), 1)\n                softmax_kernel_non_inner[grid](\n                    out,\n                    inp,\n                    M,\n                    N,\n                    K,\n                )\n            else:\n                grid = (M, 1, 1)\n                softmax_kernel_inner[grid](\n                    out,\n                    inp,\n                    M,\n                    N,\n                )\n        ctx.save_for_backward(out)\n        ctx.dim = dim\n        return out\n\n    @staticmethod\n    def backward(ctx, out_grad):\n        dim = ctx.dim\n        (out,) = ctx.saved_tensors\n        assert dim >= -out.ndim and dim < out.ndim, \"Invalid dim\"\n        dim = dim % out.ndim\n        M = 1\n        N = out.shape[dim]\n        for i in range(dim):\n            M *= out.shape[i]\n\n        out_grad = out_grad.contiguous()\n        in_grad = torch.empty_like(out)\n        K = out.numel() // M // N\n\n        with torch.cuda.device(in_grad.device):\n            if K > 1:\n                grid = lambda meta: (M, triton.cdiv(K, meta[\"TILE_K\"]), 1)\n                softmax_backward_kernel_non_inner[grid](\n                    out,\n                    out_grad,\n                    in_grad,\n                    M,\n                    N,\n                    K,\n                )\n            else:\n                grid = lambda meta: (triton.cdiv(M, meta[\"TILE_M\"]), 1, 1)\n                softmax_backward_kernel_inner[grid](\n                    out,\n                    out_grad,\n                    in_grad,\n                    M,\n                    N,\n                )\n        return in_grad, None, None\n\ndef softmax(x, dim=-1, dtype=None):\n    return Softmax.apply(x, dim, dtype)\n", "perf_candidates": [], "perf_strategy": null, "call_candidate": null, "exe_candidate": null, "temp_strategy": null, "perf_debug_num": 0, "pass_call": false, "pass_exe": false, "pass_perf": false}}