mstar.engine.resources.attn.wrappers

mstar.engine.resources.attn.wrappers#

FlashInfer utility wrappers for batched paged attention.

Provides: - FlashInferPrefillWrapper: batched prefill with paged KV cache, optional CUDA graph mode - FlashInferDecodeWrapper: batched decode with paged KV cache, optional CUDA graph mode

CUDA graph mode requires: - Static buffer pointers passed at construction (qo_indptr_buf, paged_kv_indptr_buf, etc.) - plan() updates values via .copy_() without reallocating - The same wrapper object must be used during both capture and replay

Adapted from VoxServe’s flashinfer_utils.py for our KV cache layout:

[num_layers, max_pages, 2, page_size, num_kv_heads, head_dim]

(VoxServe uses [n_pages, 2, page_size, n_heads, head_dim] without layer dim.)

Functions

check_flashinfer_head_dim(head_dim[, device])

Raise unless FlashInfer's paged attention computes head_dim exactly.

Classes

FlashInferDecodeWrapper(workspace_buffer, ...)

Batched decode attention with paged KV cache.

FlashInferMixedWrapper(decode_workspace, ...)

A step mixing one-token rows (decode) with prompt rows, on two plans.

FlashInferPrefillWrapper(workspace_buffer, ...)

Batched prefill attention with paged KV cache.
class mstar.engine.resources.attn.wrappers.FlashInferDecodeWrapper(workspace_buffer, num_qo_heads, num_kv_heads, head_dim, page_size, batch_size=None, max_num_pages=None, device=torch.device('cuda'), accelerator_graph=False, enable_nvtx=False, backend='auto')[source]#

Bases: object

Batched decode attention with paged KV cache.

Optimized for the common decode case where each request appends exactly 1 new token. Uses BatchDecodeWithPagedKVCacheWrapper.

Parameters:

  • workspace_buffer (Tensor) – FlashInfer workspace

  • num_qo_heads (int) – number of query/output heads

  • num_kv_heads (int) – number of key/value heads

  • head_dim (int) – dimension per head

  • page_size (int) – KV cache page size

  • batch_size (int | None) – required for CUDA graph mode (max requests in batch)

  • max_num_pages (int | None) – required for CUDA graph mode (max pages across all requests)

  • device (device) – torch device

  • accelerator_graph (bool) – if True, pre-allocate static buffers for graph capture

  • enable_nvtx (bool)

  • backend (str)

plan(paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len, dtype=torch.bfloat16, kv_lens=None, window_left=-1, **kwargs)[source]#

Plan decode attention and compute KV write locations.

For decode, each request appends exactly 1 token. The write location is the last page at position = last_page_len (before the append; after append it becomes last_page_len). window_left is as in the prefill wrapper’s plan.

Inputs may be on CPU; see prefill wrapper’s plan docstring.

Parameters:
run(q, kv_cache_layer)[source]#

Run planned batched decode attention.

Parameters:

  • q (Tensor) – [n_req, num_qo_heads, head_dim]

  • kv_cache_layer (Tensor) – [max_pages, 2, page_size, num_kv_heads, head_dim]

Returns:

[n_req, num_qo_heads, head_dim]

Return type:

output

class mstar.engine.resources.attn.wrappers.FlashInferMixedWrapper(decode_workspace, prefill_workspace, num_qo_heads, num_kv_heads, head_dim, page_size, batch_size, max_total_tokens, max_num_pages, max_prefill_rows, device, backend='auto')[source]#

Bases: object

A step mixing one-token rows (decode) with prompt rows, on two plans.

The packed prefill planner over every row costs per row about twice what the decode planner does, and its kernel runs the one-token rows as prefill queries. Here the decode wrapper takes every row of the bucket (a one-token row attends its own pages; any other row reads one token of the sink page, and its output is dropped), and the prefill wrapper takes at most max_prefill_rows prompt rows, their tokens moved to the front by a planned permutation. Both are captured at static shapes: the bucket’s rows and tokens, and max_prefill_rows sequences.

Every host array is staged in page-locked buffers this wrapper owns: a mixed step plans inline on the gpu thread, behind the step in flight.

Parameters:

  • decode_workspace (Tensor)

  • prefill_workspace (Tensor)

  • num_qo_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • page_size (int)

  • batch_size (int)

  • max_total_tokens (int)

  • max_num_pages (int)

  • max_prefill_rows (int)

  • device (device)

  • backend (str)

plan(qo_indptr, paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len, causal=True, dtype=torch.bfloat16, kv_lens=None, window_left=-1, **kwargs)[source]#
Parameters:
run(q, kv_cache_layer)[source]#
Parameters:
Return type:

Tensor

class mstar.engine.resources.attn.wrappers.FlashInferPrefillWrapper(workspace_buffer, num_qo_heads, num_kv_heads, head_dim, page_size, batch_size=None, max_total_tokens=None, max_num_pages=None, device=torch.device('cuda'), accelerator_graph=False, enable_nvtx=False, backend='auto')[source]#

Bases: object

Batched prefill attention with paged KV cache.

Wraps flashinfer.BatchPrefillWithPagedKVCacheWrapper with optional CUDA graph mode using static buffers. KV writes are the KVManager’s job.

Parameters:

  • workspace_buffer (Tensor) – FlashInfer workspace (256MB+ recommended)

  • num_qo_heads (int) – number of query/output heads

  • num_kv_heads (int) – number of key/value heads

  • head_dim (int) – dimension per head

  • page_size (int) – KV cache page size

  • batch_size (int | None) – required for CUDA graph mode (max requests in batch)

  • max_total_tokens (int | None) – required for CUDA graph mode (max total new tokens across batch)

  • max_num_pages (int | None) – required for CUDA graph mode (max pages across all requests)

  • device (device) – torch device

  • accelerator_graph (bool) – if True, pre-allocate static buffers for graph capture

  • enable_nvtx (bool)

  • backend (str)

plan(qo_indptr, paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len, causal=True, dtype=torch.bfloat16, kv_lens=None, window_left=-1, **kwargs)[source]#

Plan attention and compute KV write indices.

window_left is the number of preceding positions each query sees; -1 is unbounded.

In CUDA graph mode, updates static buffers via .copy_() so that the same GPU addresses are used during graph replay.

Inputs may be on CPU — that’s preferred because FlashInfer’s BatchPrefillWithPagedKVCacheWrapper.plan does indptr.to("cpu") / last_page_len.to("cpu") internally; passing GPU tensors there triggers a synchronous default-stream sync that drains the speculatively-queued next decode step. We let the inner plan consume them as CPU and async-H2D copy to the device for our own per-token bookkeeping below.

Parameters:
run(q, kv_cache_layer)[source]#

Run planned batched prefill attention.

Parameters:

  • q (Tensor) – [total_tokens, num_qo_heads, head_dim]

  • kv_cache_layer (Tensor) – [max_pages, 2, page_size, num_kv_heads, head_dim] (single layer slice of the full KV cache)

Returns:

[total_tokens, num_qo_heads, head_dim]

Return type:

output

mstar.engine.resources.attn.wrappers.check_flashinfer_head_dim(head_dim, device=None)[source]#

Raise unless FlashInfer’s paged attention computes head_dim exactly. Only CUDA runs FlashInfer, so a non-CUDA device (CPU tests driving the wrappers with a stand-in module and toy head dims) is not checked.

Parameters:
Return type:

None