mstar.engine.resources.attn.config

What a model declares about attention: its backend, its spec, its step.

Kept free of the managers and their kernels so a submodule can declare a step without pulling FlashInfer in behind it.

Classes

AttentionConfig(kv_cache, backend, ...)
AttentionSpec(resource_key, nodes, config)
AttentionStep([segments, causal])
AttnBackend(*values)
CrossAttentionConfig(kv_cache[, ...])	Cross-attention against a context written once and never extended.
CrossAttentionSpec(resource_key, nodes, config)

class mstar.engine.resources.attn.config.AttentionConfig(kv_cache: str, backend: mstar.engine.resources.attn.config.AttnBackend = <AttnBackend.FLASHINFER: 'flashinfer'>, flashinfer_backend: str = 'auto')

Bases: object

Parameters:

• kv_cache (str)

• backend (AttnBackend)

• flashinfer_backend (str)

backend: AttnBackend = 'flashinfer'

flashinfer_backend: str = 'auto'

kv_cache: str

class mstar.engine.resources.attn.config.AttentionSpec(resource_key: str, nodes: set[str], config: mstar.engine.resources.attn.config.AttentionConfig)

Bases: NodeResourceSpec

Parameters:

• resource_key (str)

• nodes (set[str])

• config (AttentionConfig)

apply_yaml_overrides(backend=None, flashinfer_backend=None)

Which kernel to run is the deployment’s call as much as the model’s — an image that cannot build FA3 pins FA2 here.

Cache geometry is not repeated here: it belongs to the KV resource this spec depends on, and is tuned under that resource’s own block.

Parameters:

• backend (str | AttnBackend | None)

• flashinfer_backend (str | None)

config: AttentionConfig

depends_on()

Keys whose specs this one builds against; the engine resolves them into EngineResourceInfo.dependencies.

Return type:

set[str]

property resource_class: type[Resource]

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

class mstar.engine.resources.attn.config.AttentionStep(segments: tuple[mstar.engine.resources.step.Segment, ...] | None = None, causal: bool = True)

Bases: ResourceStep

Parameters:

• segments (tuple[Segment, ...] | None)

• causal (bool)

causal: bool = True

class mstar.engine.resources.attn.config.AttnBackend(*values)

Bases: Enum

DENSE = 'dense'

FLASHINFER = 'flashinfer'

class mstar.engine.resources.attn.config.CrossAttentionConfig(kv_cache, query_kv_cache=None, context_label='context', backend=AttnBackend.FLASHINFER, flashinfer_backend='auto')

Bases: object

Cross-attention against a context written once and never extended.

kv_cache names the KV resource holding the encoder context; query_kv_cache names the decoder’s KV resource, whose plan defines this step’s query packing. They may be the same resource when the context shares the decoder’s head config — the context then lives in it under its own context_label. They differ when it does not, which is the usual case (an encoder’s head count rarely matches the decoder’s).

query_kv_cache=None covers the query side having no KV cache at all (nothing is cached across steps on it): the packing then comes off the cross-attention step’s own segments, one qo entry per segment in declared order.

Parameters:

• kv_cache (str)

• query_kv_cache (str | None)

• context_label (str)

• backend (AttnBackend)

• flashinfer_backend (str)

backend: AttnBackend = 'flashinfer'

context_label: str = 'context'

flashinfer_backend: str = 'auto'

kv_cache: str

query_kv_cache: str | None = None

class mstar.engine.resources.attn.config.CrossAttentionSpec(resource_key: str, nodes: set[str], config: mstar.engine.resources.attn.config.CrossAttentionConfig)

Bases: NodeResourceSpec

Parameters:

• resource_key (str)

• nodes (set[str])

• config (CrossAttentionConfig)

apply_yaml_overrides(backend=None, flashinfer_backend=None)

Which kernel to run against the context cache; see AttentionSpec.

Parameters:

• backend (str | AttnBackend | None)

• flashinfer_backend (str | None)

config: CrossAttentionConfig

depends_on()

Keys whose specs this one builds against; the engine resolves them into EngineResourceInfo.dependencies.

Return type:

set[str]

property resource_class: type[Resource]

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.