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)

Bases: AttentionManager

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

Parameters:

• kv_cache (str)

• device (device)

depends_on()

plan(step, ctx)

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

Parameters:

• step (AttentionStep)

• ctx (StepContext)

qo_indptr_buf(label='main')

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)

Parameters:

• q (Tensor)

• label (str | None)

• kv_cache_layer (Tensor | None)

• k (Tensor | None)

• v (Tensor | None)

• layer_idx (int | None)

Return type:

Tensor

select_last_hidden(hidden, label='main')

Parameters:

• hidden (Tensor)

• label (str)

Return type:

Tensor

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)

Bases: object

Parameters:

• block_table (Tensor)

• cu_q (Tensor)

• host_kv_lens (Tensor)

• max_q (int)

• max_k (int)

• causal (bool)

block_table: Tensor

causal: bool

cu_q: Tensor

host_kv_lens: Tensor

max_k: int

max_q: int