mstar.engine.resources.kv.plan

The KV plan: what KVManager.plan hands its dependents.

Attention and positions read this and nothing else of the KV resource — the packed ordering, one view per (request, label) stream, and the index tensors a paged-attention wrapper plans against. Kept apart from the manager so a consumer depends on the contract rather than on the cache.

Functions

build_paged_indptrs(segments, page_size)
group_by_plan_label(segments, combined_labels)	Segments per plan label in order of packed forward view

Classes

KVPlanOutput(cpu_indptrs, views[, cuda_indptrs])	Output of KVManager.plan for a single label
KVPlanOutputs([mapping, pre_forks, post_forks])	Plan label -> output, plus the label forks of the step it came from.
PagedIndptrs(qo_indptr, paged_kv_indptr, ...)	The four int32 index tensors a FlashInfer prefill/decode wrapper's plan consumes, built on CPU (so wrapper.plan's .to("cpu") is a no-op — see FlashInferAttentionManager.plan).
SequenceView(request_id, label, page_idxs, ...)	One (request, label) stream as this step sees it.

class mstar.engine.resources.kv.plan.KVPlanOutput(cpu_indptrs, views, cuda_indptrs=None)

Bases: object

Output of KVManager.plan for a single label

Parameters:

• cpu_indptrs (PagedIndptrs)

• views (list[SequenceView])

• cuda_indptrs (PagedIndptrs | None)

cpu_indptrs: PagedIndptrs

cuda_indptrs: PagedIndptrs | None = None

get_total_len()

property is_decode: bool

views: list[SequenceView]

class mstar.engine.resources.kv.plan.KVPlanOutputs(mapping=None, pre_forks=(), post_forks=())

Bases: dict[str, KVPlanOutput]

Plan label -> output, plus the label forks of the step it came from.

A dict subclass so consumers that only want the per-label plans keep reading it as the mapping it is. Positions need the forks too — a fork target inherits the source’s counter the way it inherits its pages — and a fork target is not necessarily a segment of the step, so it has nowhere else to ride.

Parameters:

• mapping (dict[str, KVPlanOutput] | None)

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

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

post_forks

pre_forks

class mstar.engine.resources.kv.plan.PagedIndptrs(qo_indptr, paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len)

Bases: NamedTuple

The four int32 index tensors a FlashInfer prefill/decode wrapper’s plan consumes, built on CPU (so wrapper.plan’s .to("cpu") is a no-op — see FlashInferAttentionManager.plan).

Parameters:

• qo_indptr (Tensor)

• paged_kv_indptr (Tensor)

• paged_kv_indices (Tensor)

• paged_kv_last_page_len (Tensor)

paged_kv_indices: Tensor

Alias for field number 2

paged_kv_indptr: Tensor

Alias for field number 1

paged_kv_last_page_len: Tensor

Alias for field number 3

qo_indptr: Tensor

Alias for field number 0

to_device(device)

Parameters:

device (device)

to_kwargs_dict()

class mstar.engine.resources.kv.plan.SequenceView(request_id, label, page_idxs, length, to_compute, start=0, generation=0)

Bases: NamedTuple

One (request, label) stream as this step sees it. A NamedTuple because it is built per segment per step and never mutated.

Parameters:

• request_id (str)

• label (str)

• page_idxs (list[int])

• length (int)

• to_compute (int)

• start (int)

• generation (int)

generation: int

Alias for field number 6

label: str

Alias for field number 1

last_page_len(page_size)

Parameters:

page_size (int)

Return type:

int

length: int

Alias for field number 3

page_idxs: list[int]

Alias for field number 2

request_id: str

Alias for field number 0

start: int

Alias for field number 5

to_compute: int

Alias for field number 4

mstar.engine.resources.kv.plan.build_paged_indptrs(segments, page_size)

Parameters:

• segments (list[SequenceView])

• page_size (int)

Return type:

PagedIndptrs

mstar.engine.resources.kv.plan.group_by_plan_label(segments, combined_labels)

Segments per plan label in order of packed forward view

combined plan concats source labels in label major order. standalone keeps og batch order. KV should be sole producer of this ordering and eveyrone else will read plan output of KV

NOTE: combined key with a source label with no segments in step will cause KeyError

Parameters:

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

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

Return type:

dict[str, list[Segment]]