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 — and AttentionManager.build reaches them by deferred import, so naming a backend in a spec does not load the other two.

Classes

AttentionManager()

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:
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.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