flashinfer / wrapper1b7890

import torch
import flashinfer

# group_size=6 (24 qo_heads / 4 kv_heads) is not natively supported.
# Work-around: expand KV heads from 4 to 24 (repeat_interleave x6)
# so the wrapper sees group_size=1 (MHA), which is mathematically equivalent.

_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.uint8, 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.BatchDecodeWithPagedKVCacheWrapper(workspace, kv_layout="NHD")
        _wrapper_cache[key] = wrapper
    return wrapper


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
    # kv_last_page_len may have an extra trailing element from workload capture;
    # always clamp to batch_size.
    kv_last_page_len = kv_last_page_len[:batch_size]
    group_size = num_qo_heads // num_kv_heads
    # Expand KV heads: [pages, page_size, 4, head_dim] -> [pages, page_size, 24, head_dim]
    k_exp = k_cache.repeat_interleave(group_size, dim=2)
    v_exp = v_cache.repeat_interleave(group_size, dim=2)
    expanded_kv_heads = num_qo_heads  # 24

    device = q.device
    wkey = (str(device), num_qo_heads, expanded_kv_heads, head_dim, page_size, q.dtype, k_exp.dtype)
    wrapper = _get_wrapper(wkey, device)
    state = _plan_state.get(wkey)

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

    if needs_plan:
        wrapper.plan(
            indptr=kv_indptr,
            indices=kv_indices,
            last_page_len=kv_last_page_len,
            num_qo_heads=num_qo_heads,
            num_kv_heads=expanded_kv_heads,
            head_dim=head_dim,
            page_size=page_size,
            pos_encoding_mode="NONE",
            q_data_type=q.dtype,
            kv_data_type=k_exp.dtype,
            sm_scale=sm_scale,
        )
        _plan_state[wkey] = {
            "batch_size": batch_size,
            "kv_indptr_ptr": kv_indptr.data_ptr(),
            "kv_indices_ptr": kv_indices.data_ptr(),
            "sm_scale": sm_scale,
        }

    output, lse = wrapper.run(q, (k_exp, v_exp), return_lse=True)
    return output, lse