mstar.engine.resources.attn.flashinfer_mla

mstar.engine.resources.attn.flashinfer_mla#

Multi-head Latent Attention through FlashInfer’s BatchMLAPagedAttentionWrapper.

The cache (KVLayout.MLA) holds one [ckv | kpe] latent per token; the layer writes it through the KV resource and attends with the absorbed query (q_nope @ W_UK -> kv_lora_rank wide, plus the un-rotated q_pe). The kernel’s output is in the latent space ([tokens, heads, kv_lora_rank]); the layer applies W_UV. One wrapper class serves prefill and decode (query lengths come from qo_indptr), so the slot/label bookkeeping is the same as FlashInferManager’s.

Functions

flashinfer_mla_supports(kv_lora_rank, ...)

FlashInfer's Hopper MLA kernel is specialised for the DeepSeek/Kimi latent shape (ckv 512, kpe 0 or 64); other shapes are JIT-compiled but read out of bounds.

Classes

FlashInferMLAManager(kv_cache, device, ...)

FlashInferMLAWrapper(workspace_buffer, ...)

Thin wrapper over flashinfer.mla.BatchMLAPagedAttentionWrapper.
class mstar.engine.resources.attn.flashinfer_mla.FlashInferMLAManager(kv_cache, device, dtype, kv_config, backend='auto', sm_scale=None)[source]#

Bases: AttentionManager

Parameters:
clear_preplan()[source]#
depends_on()[source]#
property force_double_buffer#

Whether this resource needs double-buffering even off the pre-plan path.

A resource that stages a step’s layout into a reused host buffer and issues a non-blocking H2D into a graph-read device buffer has a race the moment the CPU runs ahead of the GPU: plan(N+1) can overwrite the buffer before step N’s DMA has retired, and the replay attends with N+1’s data. The buffer need not be ours — FlashInfer’s plan holds one per wrapper, which is why the attention managers key theirs per slot. The main runner already double-buffers whenever a resource supports_preplan; this flag extends that to a runner that has no pre-plan path but still replays such a resource — notably the piecewise runner — and makes _exec_per_request fence before it reuses a slot, since its sub-steps all run inside one step. Off by default; a resource opts in only if it has this hazard.

kv_len_buf(label='main')[source]#

The rows’ planned kv lengths, int32 on device (static under a lease): for a context_only plan the length of each row’s stored context, which is where the row’s queries (and the entries the step appends) start.

Parameters:

label (str)

Return type:

Tensor

plan(step, ctx)[source]#

ret is immutable and opaque to runner; only gives to ctx.plan_results

Parameters:
qo_indptr_buf(label='main')[source]#
Parameters:

label (str)

Return type:

Tensor | None

run(q, label=None, kv_cache_layer=None, k=None, v=None, layer_idx=None, q_pe=None, return_lse=False)[source]#

q is the absorbed q_nope ([T, H, kv_lora_rank]); q_pe the rope-part query ([T, H, qk_rope_head_dim]). With return_lse also the natural log-sum-exp of the scores [T, H] fp32, for merging with another key segment.

Parameters:
select_last_hidden(hidden, label='main')[source]#
Parameters:
Return type:

Tensor

property supports_preplan#
property uses_kernel: bool#

False when the torch fallback serves the plans, which cannot be captured.

class mstar.engine.resources.attn.flashinfer_mla.FlashInferMLAWrapper(workspace_buffer, num_qo_heads, kv_lora_rank, qk_rope_head_dim, page_size, sm_scale, batch_size=None, max_num_pages=None, device=torch.device('cuda'), use_cuda_graph=False, backend='auto')[source]#

Bases: object

Thin wrapper over flashinfer.mla.BatchMLAPagedAttentionWrapper.

Parameters:

  • workspace_buffer (torch.Tensor) – FlashInfer float workspace.

  • num_qo_heads (int) – local query heads.

  • qk_rope_head_dim (int) – latent split (ckv / kpe widths).

  • page_size (int) – cache geometry.

  • max_num_pages (int | None) – cache geometry.

  • batch_size (int | None) – rows (requests) for CUDA-graph mode; fixed for the wrapper’s life.

  • sm_scale (float) – softmax scale of the original qk_head_dim (e.g. 192 ** -0.5).

  • kv_lora_rank (int)

  • qk_rope_head_dim

  • device (torch.device)

  • use_cuda_graph (bool)

  • backend (str)

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

Inputs may be CPU tensors (the KV plan builds them on the host).

Parameters:
run(q_nope, q_pe, kv_cache_layer, return_lse=False)[source]#

q_nope [T, H, kv_lora_rank] (absorbed), q_pe [T, H, qk_rope_head_dim], kv_cache_layer [pages, page_size, latent] -> [T, H, kv_lora_rank] (and, with return_lse, the natural log-sum-exp of the scores [T, H] fp32).

Parameters:
mstar.engine.resources.attn.flashinfer_mla.flashinfer_mla_supports(kv_lora_rank, qk_rope_head_dim)[source]#

FlashInfer’s Hopper MLA kernel is specialised for the DeepSeek/Kimi latent shape (ckv 512, kpe 0 or 64); other shapes are JIT-compiled but read out of bounds. Anything else runs the torch fallback (eager, not CUDA-graph safe): small test checkpoints, not production models.

Parameters:

  • kv_lora_rank (int)

  • qk_rope_head_dim (int)

Return type:

bool