flashinfer / wrapperea3787

import torch
import flashinfer

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


def _get_workspace(device):
    key = str(device)
    buffer = _workspace_cache.get(key)
    if buffer is None or buffer.device != device or buffer.numel() < _WORKSPACE_SIZE_BYTES:
        buffer = torch.empty(_WORKSPACE_SIZE_BYTES, dtype=torch.int8, device=device)
        _workspace_cache[key] = buffer
    return buffer


def _get_wrapper(key, device):
    wrapper = _wrapper_cache.get(key)
    if wrapper is None:
        workspace = _get_workspace(device)
        wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace)
        _wrapper_cache[key] = wrapper
    return wrapper


def run(q_nope, q_pe, ckv_cache, kpe_cache, qo_indptr, kv_indptr, kv_indices, sm_scale):
    total_q, num_qo_heads, head_dim_ckv = q_nope.shape
    _, _, head_dim_kpe = q_pe.shape
    page_size = ckv_cache.shape[1]
    len_indptr = kv_indptr.shape[0]
    num_kv_indices = kv_indices.shape[0]
    batch_size = qo_indptr.shape[0] - 1

    device = q_nope.device
    wrapper_key = (
        str(device),
        num_qo_heads,
        head_dim_ckv,
        head_dim_kpe,
        page_size,
        q_nope.dtype,
        q_pe.dtype,
        ckv_cache.dtype,
        kpe_cache.dtype,
    )

    wrapper = _get_wrapper(wrapper_key, device)
    state = _plan_state.get(wrapper_key)

    needs_plan = True
    if state is not None:
        needs_plan = (
            state.get("total_q") != total_q
            or state.get("batch_size") != batch_size
            or state.get("len_indptr") != len_indptr
            or state.get("num_kv_indices") != num_kv_indices
            or state.get("sm_scale") != sm_scale
            or state.get("qo_indptr_ptr") != qo_indptr.data_ptr()
            or state.get("kv_indptr_ptr") != kv_indptr.data_ptr()
            or state.get("kv_indices_ptr") != kv_indices.data_ptr()
        )

    if needs_plan:
        kv_len_arr = (kv_indptr[1:] - kv_indptr[:-1]).to(torch.int32)
        wrapper.plan(
            qo_indptr=qo_indptr,
            kv_indptr=kv_indptr,
            kv_indices=kv_indices,
            kv_len_arr=kv_len_arr,
            num_heads=num_qo_heads,
            head_dim_ckv=head_dim_ckv,
            head_dim_kpe=head_dim_kpe,
            page_size=page_size,
            causal=True,
            sm_scale=sm_scale,
            q_data_type=q_nope.dtype,
            kv_data_type=ckv_cache.dtype,
        )
        _plan_state[wrapper_key] = {
            "total_q": total_q,
            "batch_size": batch_size,
            "len_indptr": len_indptr,
            "num_kv_indices": num_kv_indices,
            "sm_scale": sm_scale,
            "qo_indptr_ptr": qo_indptr.data_ptr(),
            "kv_indptr_ptr": kv_indptr.data_ptr(),
            "kv_indices_ptr": kv_indices.data_ptr(),
        }

    output, lse = wrapper.run(
        q_nope,
        q_pe,
        ckv_cache,
        kpe_cache,
        return_lse=True,
    )

    return output, lse