mstar.engine.resources.attn.ragged.config

What a model declares about cacheless (ragged) attention.

Kept free of the manager and its kernels, like the other resources’ configs, so a submodule can declare a step without pulling FlashInfer in behind it.

Functions

cross_label(q_label, kv_label)

The plan a cross-attention pair runs under: q_label's spans attending kv_label's.

Classes

RaggedAttentionConfig(num_qo_heads, ...[, ...])	Varlen attention over segments packed into one forward, with no KV cache: the whole layout is this step's, and nothing carries to the next.
RaggedAttentionSpec(resource_key, nodes, config)	Cacheless self-attention within each declared span (AttentionStep).
RaggedAttnBackend(*values)
RaggedBlockCausalAttentionSpec(resource_key, ...)	Cacheless block-causal self-attention within each declared span: a token in block b of its span (block_size tokens a block, from the span's start) attends every key of blocks 0..b.
RaggedCrossAttentionSpec(resource_key, ...)	Cacheless cross-attention between two spans of each request: a forward's (q_label, kv_label) pairs (RaggedCrossAttentionStep), each run through RaggedAttentionCallable(resource, cross_label(q_label, kv_label)).
RaggedCrossAttentionStep([segments, pairs])	A ragged cross-attention step: which spans attend which.

class mstar.engine.resources.attn.ragged.config.RaggedAttentionConfig(num_qo_heads, num_kv_heads, head_dim, sm_scale=None, max_segments_per_request=1, max_tokens_per_request=None, backend=RaggedAttnBackend.AUTO, flashinfer_backend='auto', dtype=None)

Bases: object

Varlen attention over segments packed into one forward, with no KV cache: the whole layout is this step’s, and nothing carries to the next. The config of both RaggedAttentionSpec (self-attention within each span) and RaggedCrossAttentionSpec (one span attending another).

Head counts are pre-sharding; the engine narrows them to the rank’s slice at build, as it does for a KVConfig.

Parameters:

• num_qo_heads (int)

• num_kv_heads (int)

• head_dim (int)

• sm_scale (float | None)

• max_segments_per_request (int)

• max_tokens_per_request (int | None)

• backend (RaggedAttnBackend)

• flashinfer_backend (str)

• dtype (dtype | None)

backend: RaggedAttnBackend = 'auto'

dtype: dtype | None = None

flashinfer_backend: str = 'auto'

head_dim: int

max_segments_for(bs)

Parameters:

bs (int)

Return type:

int

max_segments_per_request: int = 1

max_tokens_for(bs)

Parameters:

bs (int)

Return type:

int | None

max_tokens_per_request: int | None = None

num_kv_heads: int

num_qo_heads: int

shard(num_shards)

Narrow the head counts to one rank’s slice; see KVConfig.shard.

Idempotent, so a rebuild (or a second manager over one config) is free.

Parameters:

num_shards (int)

Return type:

None

sm_scale: float | None = None

class mstar.engine.resources.attn.ragged.config.RaggedAttentionSpec(resource_key, nodes, config)

Bases: _RaggedSpecBase

Cacheless self-attention within each declared span (AttentionStep).

Parameters:

• resource_key (str)

• nodes (set[str])

• config (RaggedAttentionConfig)

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.ragged.config.RaggedAttnBackend(*values)

Bases: Enum

AUTO = 'auto'

FLASHINFER = 'flashinfer'

class mstar.engine.resources.attn.ragged.config.RaggedBlockCausalAttentionSpec(resource_key, nodes, config, block_size)

Bases: _RaggedSpecBase

Cacheless block-causal self-attention within each declared span: a token in block b of its span (block_size tokens a block, from the span’s start) attends every key of blocks 0..b. Declared with AttentionStep(causal=False), since the block structure is the whole mask.

Its own kind, not a flag on RaggedAttentionSpec: it runs a different kernel (paged prefill over one-token pages, see attn.ragged.block_causal).

Parameters:

• resource_key (str)

• nodes (set[str])

• config (RaggedAttentionConfig)

• block_size (int)

block_size: int

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.ragged.config.RaggedCrossAttentionSpec(resource_key, nodes, config)

Bases: _RaggedSpecBase

Cacheless cross-attention between two spans of each request: a forward’s (q_label, kv_label) pairs (RaggedCrossAttentionStep), each run through RaggedAttentionCallable(resource, cross_label(q_label, kv_label)).

Separate from RaggedAttentionSpec so a resource is one kind of attention: a model with both declares one of each per head geometry.

Parameters:

• resource_key (str)

• nodes (set[str])

• config (RaggedAttentionConfig)

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.ragged.config.RaggedCrossAttentionStep(segments=None, pairs=())

Bases: ResourceStep

A ragged cross-attention step: which spans attend which.

Each label in segments is one span per request (video tokens, text tokens, …). Each (q_label, kv_label) in pairs is one attention: every request’s q_label span attends that request’s kv_label span, so both labels must be declared for the same requests in the same order. Never causal.

Parameters:

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

• pairs (tuple[tuple[str, str], ...])

pairs: tuple[tuple[str, str], ...] = ()

mstar.engine.resources.attn.ragged.config.cross_label(q_label, kv_label)

The plan a cross-attention pair runs under: q_label’s spans attending kv_label’s.

Parameters:

• q_label (str)

• kv_label (str)

Return type:

str