flashinfer / wrapper463e16

import torch
import flashinfer

# GQA group_size=2 (32 qo_heads / 16 kv_heads), Gemma 3 27B at TP=1.

_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.BatchPrefillWithRaggedKVCacheWrapper(_get_workspace(device), kv_layout="NHD")
        _wrapper_cache[key] = w
    return w


def run(q, k, v, qo_indptr, kv_indptr, sm_scale):
    total_q, num_qo_heads, head_dim = q.shape
    total_kv, num_kv_heads, _ = k.shape
    device = q.device
    wkey = (str(device), num_qo_heads, num_kv_heads, head_dim, q.dtype, k.dtype)
    wrapper = _get_wrapper(wkey, device)
    state = _plan_state.get(wkey)
    needs_plan = state is None or state["total_q"] != total_q or state["qo_ptr"] != qo_indptr.data_ptr()
    if needs_plan:
        wrapper.plan(
            qo_indptr=qo_indptr,
            kv_indptr=kv_indptr,
            num_qo_heads=num_qo_heads,
            num_kv_heads=num_kv_heads,
            head_dim_qk=head_dim,
            causal=True,
            sm_scale=float(sm_scale),
            q_data_type=q.dtype,
            kv_data_type=k.dtype,
        )
        _plan_state[wkey] = {"total_q": total_q, "qo_ptr": qo_indptr.data_ptr()}
    output, lse = wrapper.run(q, k, v, return_lse=True)
    return output, lse