mstar.engine.resources.kv.bounded.config

mstar.engine.resources.kv.bounded.config#

What a model declares about a bounded KV cache.

A request’s cache is a logical stream: an optional read-only source prefix (a voice prompt’s keys, say, shared by every request that uses it) followed by the tokens the request writes. Retention keeps the stream’s first sink entries and its last window (StreamingLLM’s attention sinks plus a sliding window), so a request’s footprint is fixed and known when it starts: each slot holds only the written part of the sink and a ring for the window, and source entries are read where they live.

Classes

BoundedKVConfig(num_layers, num_heads, ...)

BoundedKVSpec(resource_key, nodes, config)

BoundedKVStep([segments, positions])

Each segment's span is what its request writes this step (0 reads only).

SinkWindow(sink, window[, window_includes_step])

Keep the stream's first sink and last window entries.

StreamPosition(source_len, written[, source])

Where a request's stream stands when a step starts.
class mstar.engine.resources.kv.bounded.config.BoundedKVConfig(num_layers: int, num_heads: int, head_dim: int, rows_per_request: int, max_source_len: int, retention: collections.abc.Callable[[int], mstar.engine.resources.kv.bounded.config.SinkWindow] | mstar.engine.resources.kv.bounded.config.SinkWindow, max_slots: int = 32, min_source_len: int = 0, dtype: torch.dtype | None = torch.float32, reverse_step_order: bool = False, query_window: int | None = None)[source]#

Bases: object

Parameters:
dtype: dtype | None = torch.float32#
head_dim: int#
property max_retained: int#

The most entries a request retains, over every allowed source length.

max_slots: int = 32#
max_source_len: int#
min_source_len: int = 0#
num_heads: int#
num_layers: int#
query_window: int | None = None#
retention: Callable[[int], SinkWindow] | SinkWindow#
retention_for(source_len)[source]#
Parameters:

source_len (int)

Return type:

SinkWindow

reverse_step_order: bool = False#
rows_per_request: int#
property sink_capacity: int#

Written tokens a slot’s sink region holds (the sink past the source).

property slot_tokens: int#
property window_capacity: int#
class mstar.engine.resources.kv.bounded.config.BoundedKVSpec(resource_key: str, nodes: set[str], config: mstar.engine.resources.kv.bounded.config.BoundedKVConfig)[source]#

Bases: NodeResourceSpec

Parameters:
apply_yaml_overrides(max_slots=None)[source]#

Patch declared parameters from this resource’s YAML block.

The model declares shapes that suit the model; a deployment tunes what suits the box it runs on. The block is scoped to this spec’s resource_key (see apply_yaml_overrides below), so an unrecognized key here is a typo, not another resource’s setting — subclasses name exactly what they accept and let the rest raise.

TODO: generalize. Matching YAML keys against the spec’s (and its config’s) dataclass fields would remove these per-resource implementations, at the cost of silently accepting anything named alike.

Parameters:

max_slots (int | None)

config: BoundedKVConfig#
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.bounded.config.BoundedKVStep(segments=None, positions=<factory>)[source]#

Bases: ResourceStep

Each segment’s span is what its request writes this step (0 reads only).

The positions come from the model, not the resource, so a plan never waits on the previous step’s commit and planning the same step twice (a piecewise region inside a forward that declared it too) is harmless.

Parameters:
positions: Mapping[str, StreamPosition]#
class mstar.engine.resources.kv.bounded.config.SinkWindow(sink, window, window_includes_step=False)[source]#

Bases: NamedTuple

Keep the stream’s first sink and last window entries. With window_includes_step the window counts the step’s own tokens, so a step that writes span attends only the last window - span retained ones.

Parameters:

  • sink (int)

  • window (int)

  • window_includes_step (bool)

sink: int#

Alias for field number 0

window: int#

Alias for field number 1

window_includes_step: bool#

Alias for field number 2

class mstar.engine.resources.kv.bounded.config.StreamPosition(source_len, written, source=0)[source]#

Bases: NamedTuple

Where a request’s stream stands when a step starts.

Parameters:

  • source_len (int)

  • written (int)

  • source (int)

source: int#

Alias for field number 2

source_len: int#

Alias for field number 0

written: int#

Alias for field number 1