flashinfer_wrapper_bcbabf

import torch
import flashinfer

# GQA group_size=5 (20 qo_heads / 4 kv_heads) is not a power-of-2 and is
# unsupported by FlashInfer's BatchDecodeWithPagedKVCacheWrapper.  Work-around:
# expand KV heads from 4 to 20 (repeat_interleave x5) so group_size=1 (MHA),
# which is mathematically equivalent.  We use BatchPrefillWithPagedKVCacheWrapper
# with causal=False, treating each decode step as a 1-token prefill per sequence.

_WORKSPACE_SIZE_BYTES = 128 * 1024 * 1024
_workspace_cache = {}
_wrapper_cache = {}
_plan_state = {}


def _get_workspace(device):
    key = str(device)
    buf = _workspace_cache.get(key)
    if buf is None:
        buf = torch.empty(_WORKSPACE_SIZE_BYTES, dtype=torch.uint8, device=device)
        _workspace_cache[key] = buf
    return buf


def _get_wrapper(key, device):
    w = _wrapper_cache.get(key)
    if w is None:
        w = flashinfer.BatchPrefillWithPagedKVCacheWrapper(_get_workspace(device), kv_layout="NHD")
        _wrapper_cache[key] = w
    return w


def run(q, k_cache, v_cache, kv_indptr, kv_indices, kv_last_page_len, sm_scale):
    batch_size, num_qo_heads, head_dim = q.shape
    _, page_size, num_kv_heads, _ = k_cache.shape
    device = q.device
    group_size = num_qo_heads // num_kv_heads  # 5
    # Expand KV heads: [num_pages, page_size, 4, 128] -> [num_pages, page_size, 20, 128]
    k_exp = k_cache.repeat_interleave(group_size, dim=2)
    v_exp = v_cache.repeat_interleave(group_size, dim=2)
    paged_kv = torch.stack([k_exp, v_exp], dim=1)  # [num_pages, 2, page_size, 20, 128]
    expanded_heads = num_qo_heads  # 20
    wkey = (str(device), num_qo_heads, expanded_heads, head_dim, page_size, q.dtype, k_cache.dtype)
    wrapper = _get_wrapper(wkey, device)
    state = _plan_state.get(wkey)
    needs_plan = (
        state is None
        or state["batch_size"] != batch_size
        or state["kv_ptr"] != kv_indptr.data_ptr()
        or state["last_page_ptr"] != kv_last_page_len.data_ptr()
    )
    if needs_plan:
        qo_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device=device)
        wrapper.plan(
            qo_indptr=qo_indptr,
            paged_kv_indptr=kv_indptr,
            paged_kv_indices=kv_indices,
            paged_kv_last_page_len=kv_last_page_len,
            num_qo_heads=num_qo_heads,
            num_kv_heads=expanded_heads,
            head_dim_qk=head_dim,
            page_size=page_size,
            causal=False,
            sm_scale=float(sm_scale),
            q_data_type=q.dtype,
            kv_data_type=k_cache.dtype,
        )
        _plan_state[wkey] = {
            "batch_size": batch_size,
            "kv_ptr": kv_indptr.data_ptr(),
            "last_page_ptr": kv_last_page_len.data_ptr(),
        }
    output, lse = wrapper.run(q, paged_kv, return_lse=True)
    return output, lse