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.

Functions

refuse_undeclared_labels(model, declared, ...)	The model is at fault and not the client, so this raises RuntimeError, which fails the request with a 500.
without_layouts(configs)	Only the KV manager's ingest reads the layouts and caps, so they are dropped before the configs go out with every frame between workers.

Classes

KVConfig(*, num_layers, num_kv_heads, head_dim)	The model geometry every KV storage strategy needs, and nothing else.
KVLayout(*values)
KVReqConfig(needed_labels, ...)
KVSpec(resource_key, nodes, config)
KVStep(segments, ...] | None = None, commit, ...)
PagedKVConfig(*, num_layers, num_kv_heads, ...)	Fixed-size pages, appended to as a sequence grows.
PrefixSpan(length, advance, walk[, scratch, ...])	One write to a keyed stream, as the preprocess worker laid it out.
RetentionPolicy(context_budget[, ...])	FIFO retention for a stream that keeps committing (windowed / rolling generation): once the committed tokens behind protected_prefix exceed context_budget, the oldest unprotected pages are released.
RingKVConfig(*, num_layers, num_kv_heads, ...)	A fixed horizon of frame slots per layer, overwritten in place.
RingKVLayerConfig(ring_frames, ring_buckets, ...)	One layer's ring geometry.
RingKVStep([segments])	One ring clock per request in the step's batch.
StreamLayout([spans, growth, scratch, ...])	LabelLayout as the KV manager holds it.

class mstar.engine.resources.kv.config.KVConfig(*, num_layers, num_kv_heads, head_dim, num_qo_heads=None)

Bases: ABC

The model geometry every KV storage strategy needs, and nothing else.

What a cache holds (layers, heads, head dim) is a checkpoint fact and lives here; how it is stored — paged or ringed, is storage policy and lives in a subclass. KVSpec dispatches on which one it was handed.

kw_only is load-bearing, not style: num_qo_heads carries a default here while both subclasses add required fields, which is an invalid positional dataclass field order.

Parameters:

• num_layers (int)

• num_kv_heads (int)

• head_dim (int)

• num_qo_heads (int | None)

abstractmethod apply_yaml_overrides(**kwargs)

Patch this deployment’s tunables; see NodeResourceSpec.

Abstract so each storage policy names exactly what it accepts and lets the rest raise — a paged key against a ring config is a typo, and the spec’s contract is that a typo is loud.

Return type:

None

head_dim: int

num_kv_heads: int

num_layers: int

num_qo_heads: int | None = None

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 session 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

MLA = 'MLA'

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>, publish_labels_per_node_walk: dict[tuple[str, str], list[str]] | None = None, final_publish_labels_per_node_walk: dict[tuple[str, str], list[str]] | None = None, prefix_keys: dict[str, list[bytes]] | None = None, prefix_tail: dict[str, list[int]] | None = None, prefix_decode: dict[str, str] | None = None, prefix_cache: bool = True, kv_layout: dict[str, 'StreamLayout'] | None = None, loop_caps: dict[str, int] | None = None)

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]])

• publish_labels_per_node_walk (dict[tuple[str, str], list[str]] | None)

• final_publish_labels_per_node_walk (dict[tuple[str, str], list[str]] | None)

• prefix_keys (dict[str, list[bytes]] | None)

• prefix_tail (dict[str, list[int]] | None)

• prefix_decode (dict[str, str] | None)

• prefix_cache (bool)

• kv_layout (dict[str, StreamLayout] | None)

• loop_caps (dict[str, int] | None)

apply_conductor_config(prefix_keys=None, prefix_tail=None, prefix_decode=None, prefix_cache=None, kv_layout=None, loop_caps=None, **kwargs)

Parameters:

• prefix_keys (dict[str, list[bytes]] | None)

• prefix_tail (dict[str, list[int]] | None)

• prefix_decode (dict[str, str] | None)

• prefix_cache (bool | None)

• kv_layout (dict[str, list] | None)

• loop_caps (Mapping[str, int] | None)

final_publish_labels_per_node_walk: dict[tuple[str, str], list[str]] | None = None

get_labels(node, walk)

Parameters:

• node (str)

• walk (str)

get_publish_labels(node, walk, available, *, final=False)

Parameters:

• node (str | None)

• walk (str | None)

• available (list[str])

• final (bool)

Return type:

list[str]

kv_layout: dict[str, StreamLayout] | None = None

loop_caps: dict[str, int] | None = None

loop_room(spec=None)

The most iterations each loop can run before the request outgrows the resource built from spec. A loop the resource doesn’t bound is left out.

Parameters:

spec (KVSpec | None)

Return type:

dict[str, int]

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]]

opened_labels()

Return type:

set[str]

prefix_cache: bool = True

prefix_decode: dict[str, str] | None = None

prefix_keys: dict[str, list[bytes]] | None = None

prefix_tail: dict[str, list[int]] | None = None

publish_labels_per_node_walk: dict[tuple[str, str], list[str]] | None = None

validate(spec=None)

ValueError (a 400) when a label would hold more positions at the request’s caps than its pool’s max_seq_len. A label under retention is left out, since Cosmos3, the one model that declares one, computes its positions in each forward.

Parameters:

spec (KVSpec | None)

Return type:

None

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(**kwargs)

Forwarded to the config: which keys are legal is a property of the storage strategy, so the config answers for them.

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], ...] = (), fork_rids: frozenset | None = None, sink: frozenset[tuple[str, str]] = frozenset(), retention: collections.abc.Mapping[tuple[str, str], mstar.engine.resources.kv.config.RetentionPolicy] = <factory>)

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], ...])

• fork_rids (frozenset | None)

• sink (frozenset[tuple[str, str]])

• retention (Mapping[tuple[str, str], RetentionPolicy])

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

commit: bool = True

fork_rids: frozenset | None = None

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

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

retention: Mapping[tuple[str, str], RetentionPolicy]

sink: frozenset[tuple[str, str]] = frozenset({})

class mstar.engine.resources.kv.config.PagedKVConfig(*, num_layers, num_kv_heads, head_dim, num_qo_heads=None, max_seq_len, max_num_pages=2048, page_size=128, layout=KVLayout.NHD, cpu_offload_pages=0, prefix_cache_salt='', prefix_cache=True, prefix_cache_tp=False, kv_lora_rank=None, qk_rope_head_dim=None)

Bases: KVConfig

Fixed-size pages, appended to as a sequence grows. The default.

Parameters:

• num_layers (int)

• num_kv_heads (int)

• head_dim (int)

• num_qo_heads (int | None)

• max_seq_len (int)

• max_num_pages (int)

• page_size (int)

• layout (KVLayout)

• cpu_offload_pages (int)

• prefix_cache_salt (str)

• prefix_cache (bool)

• prefix_cache_tp (bool)

• kv_lora_rank (int | None)

• qk_rope_head_dim (int | None)

apply_yaml_overrides(max_num_pages=None, page_size=None, max_seq_len=None, cpu_offload_pages=None, prefix_cache_salt=None, prefix_cache=None, prefix_cache_tp=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)

• prefix_cache_salt (str | None)

• prefix_cache (bool | None)

• prefix_cache_tp (bool | None)

Return type:

None

cpu_offload_pages: int = 0

kv_lora_rank: int | None = None

property latent_dim: int

Per-token cache width under MLA.

layout: KVLayout = 'NHD'

max_num_pages: int = 2048

max_seq_len: int

page_size: int = 128

prefix_cache: bool = True

prefix_cache_salt: str = ''

prefix_cache_tp: bool = False

qk_rope_head_dim: int | None = None

qo_heads_given: bool = False

class mstar.engine.resources.kv.config.PrefixSpan(length, advance, walk, scratch=0, digest=None, end_key=None)

Bases: object

One write to a keyed stream, as the preprocess worker laid it out.

length is how many slots the write fills and advance is how far it moves the position counter. digest is an item’s digest, and None for ids. end_key is set for an item that ends inside a page, and keys that page up to the item’s end.

Parameters:

• length (int)

• advance (int)

• walk (str)

• scratch (int)

• digest (bytes | None)

• end_key (bytes | None)

advance: int

digest: bytes | None = None

end_key: bytes | None = None

length: int

scratch: int = 0

walk: str

class mstar.engine.resources.kv.config.RetentionPolicy(context_budget, protected_prefix=0)

Bases: object

FIFO retention for a stream that keeps committing (windowed / rolling generation): once the committed tokens behind protected_prefix exceed context_budget, the oldest unprotected pages are released.

A committing step declares it for its stream on KVStep.retention and KVManager.commit applies it, so the release happens between steps as far as every planner is concerned: a pre-plan of the next step gates on this commit and sees the compacted stream. Whole pages only (the page straddling the prefix boundary and a partial tail page stay), so the realized context can run over the budget by up to a page; the excess is re-offered at the next commit. protected_prefix tokens at the head (a text prompt, say) are never released.

Parameters:

• context_budget (int)

• protected_prefix (int)

context_budget: int

protected_prefix: int = 0

class mstar.engine.resources.kv.config.RingKVConfig(*, num_layers, num_kv_heads, head_dim, num_qo_heads=None, tokens_per_frame, layers, num_sessions=1)

Bases: KVConfig

A fixed horizon of frame slots per layer, overwritten in place.

Every layer’s storage is allocated once and reused for the life of the process, and a write to an occupied slot is the intended behaviour.

Parameters:

• num_layers (int)

• num_kv_heads (int)

• head_dim (int)

• num_qo_heads (int | None)

• tokens_per_frame (int)

• layers (tuple[RingKVLayerConfig, ...])

• num_sessions (int)

apply_yaml_overrides(num_sessions=None, **kwargs)

num_sessions only. Nothing else here is a deployment knob.

Parameters:

num_sessions (int | None)

Return type:

None

layers: tuple[RingKVLayerConfig, ...]

num_sessions: int = 1

tokens_per_frame: int

property total_sessions: int

Resident sessions plus one shared scratch session for padding rows.

A replay padded to its capture bucket parks the dummy tail on this session (see RingKVManager.plan); it is never handed to a request, so resident capacity stays num_sessions and the deployment knob keeps its meaning. The ring buffer and the flex mask both size on this count so the padding session is a real, addressable span.

class mstar.engine.resources.kv.config.RingKVLayerConfig(ring_frames, ring_buckets, pinned_dilation)

Bases: object

One layer’s ring geometry. Layers can hold frames at different strides.

Parameters:

• ring_frames (int)

• ring_buckets (int)

• pinned_dilation (int)

pinned_dilation: int

ring_buckets: int

ring_frames: int

class mstar.engine.resources.kv.config.RingKVStep(segments=None, *, frames)

Bases: ResourceStep

One ring clock per request in the step’s batch.

Parameters:

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

• frames (tuple[tuple[str, int], ...])

frames: tuple[tuple[str, int], ...]

class mstar.engine.resources.kv.config.StreamLayout(spans=(), growth=(), scratch=0, retention=None, keyed_spans=0)

Bases: object

LabelLayout as the KV manager holds it.

Parameters:

• spans (tuple[PrefixSpan, ...])

• growth (tuple[tuple[str, int], ...])

• scratch (int)

• retention (RetentionPolicy | None)

• keyed_spans (int)

growth: tuple[tuple[str, int], ...] = ()

keyed_spans: int = 0

pages(caps, page_size)

The most pages the stream can hold, where caps gives each growth loop’s iterations.

Under retention the stream holds a window. KVManager._apply_retention leaves less than a page over the budget, so the window peaks at the protected prefix plus the budget, or at the prefix rounded up to a page if that is larger. The bound adds a page of overrun and the largest step past the prefix to that peak.

Parameters:

• caps (Mapping[str, int])

• page_size (int)

Return type:

int

positions(caps)

Parameters:

caps (Mapping[str, int])

Return type:

int

retention: RetentionPolicy | None = None

scratch: int = 0

slots(caps)

Parameters:

caps (Mapping[str, int])

Return type:

int

spans: tuple[PrefixSpan, ...] = ()

mstar.engine.resources.kv.config.refuse_undeclared_labels(model, declared, kv_layout, configs)

The model is at fault and not the client, so this raises RuntimeError, which fails the request with a 500.

Parameters:

• model (str)

• declared (Mapping[str, tuple[str, ...]])

• kv_layout (Mapping[str, Mapping])

• configs (Mapping[str, ResourceReqConfig])

Return type:

None

mstar.engine.resources.kv.config.without_layouts(configs)

Only the KV manager’s ingest reads the layouts and caps, so they are dropped before the configs go out with every frame between workers.

Parameters:

configs (Mapping[str, ResourceReqConfig])

Return type:

dict[str, ResourceReqConfig]