mstar.engine.resources.attn.ragged.config

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.

Classes

RaggedAttentionConfig(num_qo_heads, ...[, ...])

Varlen self-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)

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, flashinfer_backend='auto')[source]#

Bases: object

Varlen self-attention over segments packed into one forward, with no KV cache: the whole layout is this step’s, and nothing carries to the next.

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)

  • flashinfer_backend (str)

flashinfer_backend: str = 'auto'#
head_dim: int#
max_segments_for(bs)[source]#
Parameters:

bs (int)

Return type:

int

max_segments_per_request: int = 1#
max_tokens_for(bs)[source]#
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)[source]#

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: str, nodes: set[str], config: mstar.engine.resources.attn.ragged.config.RaggedAttentionConfig)[source]#

Bases: NodeResourceSpec

Parameters:
apply_yaml_overrides(flashinfer_backend=None, max_segments_per_request=None, max_tokens_per_request=None)[source]#

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

The two ceilings are here rather than on the model because they size CUDA-graph buckets, which is a deployment’s memory/coverage trade.

Parameters:

  • flashinfer_backend (str | None)

  • max_segments_per_request (int | None)

  • max_tokens_per_request (int | None)

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.