mstar.engine.resources.kv.config

What a model declares about a KV cache: its shape, its spec, its step.

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

Classes

KVConfig(num_layers, num_kv_heads, head_dim, ...)
KVLayout(*values)
KVReqConfig(needed_labels, ...)
KVSpec(resource_key, nodes, config)
KVStep(segments, ...] | None = None, commit, ...)

class mstar.engine.resources.kv.config.KVConfig(num_layers: int, num_kv_heads: int, head_dim: int, max_seq_len: int, max_num_pages: int = 2048, page_size: int = 128, num_qo_heads: int = None, layout: mstar.engine.resources.kv.config.KVLayout = <KVLayout.NHD: 'NHD'>, cpu_offload_pages: int = 0)

Bases: object

Parameters:

• num_layers (int)

• num_kv_heads (int)

• head_dim (int)

• max_seq_len (int)

• max_num_pages (int)

• page_size (int)

• num_qo_heads (int)

• layout (KVLayout)

• cpu_offload_pages (int)

cpu_offload_pages: int = 0

head_dim: int

layout: KVLayout = 'NHD'

max_num_pages: int = 2048

max_seq_len: int

num_kv_heads: int

num_layers: int

num_qo_heads: int = None

page_size: int = 128

shard(num_shards)

Narrow the head counts to one rank’s slice.

Idempotent because one KVConfig is shared by the KV resource and the attention resources planned against it, and each shards on construction. num_shards is the instance world size (tp * sp): Ulysses SP all-to-alls heads, so attention runs at head-degree tp*sp.

Parameters:

num_shards (int)

Return type:

None

class mstar.engine.resources.kv.config.KVLayout(*values)

Bases: Enum

NHD = 'NHD'

class mstar.engine.resources.kv.config.KVReqConfig(needed_labels: list[str] | None = None, needed_labels_per_node: dict[str, list[str]] = <factory>, needed_labels_per_node_walk: dict[tuple[str, str], list[str]] = <factory>)

Bases: ResourceReqConfig

Parameters:

• needed_labels (list[str] | None)

• needed_labels_per_node (dict[str, list[str]])

• needed_labels_per_node_walk (dict[tuple[str, str], list[str]])

get_labels(node, walk)

Parameters:

• node (str)

• walk (str)

needed_labels: list[str] | None = None

needed_labels_per_node: dict[str, list[str]]

needed_labels_per_node_walk: dict[tuple[str, str], list[str]]

class mstar.engine.resources.kv.config.KVSpec(resource_key: str, nodes: set[str], config: mstar.engine.resources.kv.config.KVConfig)

Bases: NodeResourceSpec

Parameters:

• resource_key (str)

• nodes (set[str])

• config (KVConfig)

apply_yaml_overrides(max_num_pages=None, page_size=None, max_seq_len=None, cpu_offload_pages=None)

How much cache this deployment gets, and how it is cut up.

Parameters:

• max_num_pages (int | None)

• page_size (int | None)

• max_seq_len (int | None)

• cpu_offload_pages (int | None)

config: KVConfig

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.kv.config.KVStep(segments: tuple[mstar.engine.resources.step.Segment, ...] | None = None, commit: bool = True, combined_labels: dict[tuple[str, ...], str] = <factory>, pre_forks: tuple[tuple[str, str], ...] = (), post_forks: tuple[tuple[str, str], ...] = ())

Bases: ResourceStep

Parameters:

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

• commit (bool)

• combined_labels (dict[tuple[str, ...], str])

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

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

combined_labels: dict[tuple[str, ...], str]

commit: bool = True

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

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