mstar.engine.resources.attn.base

mstar.engine.resources.attn.base#

The attention resource’s shared machinery: the spec-time factory, the custom op every backend attends through, and the workspace pool.

The backends themselves live beside this — flashinfer, cross, dense, flex — and AttentionManager.build reaches them by deferred import, so naming a backend in a spec does not load the others.

Functions

captured_page_capacity(num_rows, num_tokens, ...)

Page ids a captured wrapper's index buffer holds for a bucket of num_rows rows and num_tokens tokens.

resolve_attn_backend(backend, device)

The concrete backend for backend on device.

Classes

AttentionManager()

EagerSlotKey(label, slot[, is_decode])

One eager wrapper — the counterpart of CGSlotKey, slotted for the same reason (FlashInfer's plan stages into a pinned buffer per wrapper).

PlanCacheKey(q_seq_lens, page_indices, ...)

Fingerprint of a wrapper plan call's inputs.

WorkspacePool(device)

FlashInfer workspace buffers, one per (plan label, cg slot).
class mstar.engine.resources.attn.base.AttentionManager[source]#

Bases: AttentionResource

classmethod build(spec, info)[source]#
Parameters:
fingerprint()[source]#

What this resource contributes to the prefix cache’s root.

Everything a stored page depends on beyond its own tokens. None leaves the resource out of the root, which is right only where it cannot change what a page holds.

Return type:

bytes

prefix_skip_safe: bool = True#
property requires_kv_write: bool#

Whether a layer must write this step’s K/V through the KV resource before calling run.

False only for backends that take the fresh K/V straight into the kernel (the dense one). A layer therefore reads

if self.attn.requires_kv_write:

self.kv.write_kv(k, v, layer_idx=i, label=label)

out = self.attn.run(q, label, kv.layer_view(i), k=k, v=v, layer_idx=i)

and stays correct whichever backend the spec named.

class mstar.engine.resources.attn.base.EagerSlotKey(label, slot, is_decode=False)[source]#

Bases: NamedTuple

One eager wrapper — the counterpart of CGSlotKey, slotted for the same reason (FlashInfer’s plan stages into a pinned buffer per wrapper).

Parameters:
is_decode: bool#

Alias for field number 2

label: str#

Alias for field number 0

slot: int#

Alias for field number 1

class mstar.engine.resources.attn.base.PlanCacheKey(q_seq_lens, page_indices, last_page_lens)[source]#

Bases: NamedTuple

Fingerprint of a wrapper plan call’s inputs. When it is unchanged between steps the re-plan is skippable.

Parameters:
last_page_lens: tuple#

Alias for field number 2

page_indices: tuple#

Alias for field number 1

q_seq_lens: tuple#

Alias for field number 0

class mstar.engine.resources.attn.base.WorkspacePool(device)[source]#

Bases: object

FlashInfer workspace buffers, one per (plan label, cg slot).

Persistent: a wrapper is planned against the buffer it was built with, so these outlive any single step.

Parameters:

device (device)

get(label, cg_slot=None)[source]#
Parameters:

  • label (str)

  • cg_slot (int | None)

Return type:

Tensor

mstar.engine.resources.attn.base.captured_page_capacity(num_rows, num_tokens, kv_config)[source]#

Page ids a captured wrapper’s index buffer holds for a bucket of num_rows rows and num_tokens tokens.

Per row, not per pool: rows that matched one cached prefix all name its pages, so a step can name more page ids than the pool holds. And at least the bucket’s own pages: a capture’s rows name a SINK_PAGE per page of span, whatever the pool’s size.

Parameters:
Return type:

int

mstar.engine.resources.attn.base.resolve_attn_backend(backend, device)[source]#

The concrete backend for backend on device.

AUTO becomes the platform’s default, FlashInfer on every platform for now. A backend named in the spec or YAML is returned as is.

Parameters:
Return type:

AttnBackend