mstar.engine.resources.attn.xpu

mstar.engine.resources.attn.xpu#

Paged attention through vllm-xpu-kernels.

Classes

XPUPagedAttentionManager(kv_cache, device, ...)

Paged KV attention backed by vllm-xpu-kernels FlashAttention.

XPUPagedPlan(block_table, cu_q, ...)

class mstar.engine.resources.attn.xpu.XPUPagedAttentionManager(kv_cache, device, kv_config)[source]#

Bases: AttentionManager

Paged KV attention backed by vllm-xpu-kernels FlashAttention.

Parameters:
build_accelerator_graph_buffers(slots, max_bs, max_seq_len)[source]#

Size whatever the captured replays will read.

Called once per runner that captures against this node — the whole forward’s, and one per piecewise region — so it must tolerate repeated calls: grow to the largest shape asked for, never clobber what an earlier call already sized.

Parameters:
Return type:

None

clear_preplan()[source]#
depends_on()[source]#
plan(step, ctx)[source]#

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

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

Return this eager step’s query indptr.

Unlike CUDA-graph FlashInfer buffers, this tensor is newly allocated by every plan and therefore has no address-stability guarantee.

Parameters:

label (str)

Return type:

Tensor | None

run(q, label=None, kv_cache_layer=None, k=None, v=None, layer_idx=None)[source]#
Parameters:
Return type:

Tensor

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

Tensor

property supports_preplan#
class mstar.engine.resources.attn.xpu.XPUPagedPlan(block_table: torch.Tensor, cu_q: torch.Tensor, host_kv_lens: torch.Tensor, max_q: int, max_k: int, causal: bool)[source]#

Bases: object

Parameters:
block_table: Tensor#
causal: bool#
cu_q: Tensor#
host_kv_lens: Tensor#
max_k: int#
max_q: int#