{"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_debug_tmp/tmp/gen/lightning_attention.py_gen_triton_code_164920.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": 4, "pass_call": false, "pass_exe": false, "pass_perf": false}}