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[, ring])
context_only_views(views, query_spans, page_size)	The views of a step's streams without the tokens the step appends: length shrinks by to_compute (the pages beyond it dropped) and to_compute becomes the row's query count.
group_by_plan_label(segments, combined_labels)	Segments per plan label in order of packed forward view

Classes

HostLens	A host int32 row whose iteration yields numpy scalars.
KVPlanOutput(cpu_indptrs[, views, ...])	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).
PinnedIndexRing([numel, depth, per_plan])	A ring of page-locked int32 host buffers for the index arrays a plan hands FlashInfer, so its copy_(non_blocking=True) calls are really asynchronous (from pageable memory the driver stages each one).
SequenceView(request_id, label, page_idxs, ...)	One (request, label) stream as this step sees it.

class mstar.engine.resources.kv.plan.HostLens

Bases: Tensor

A host int32 row whose iteration yields numpy scalars.

FlashInfer’s decode plan takes the longest row with the Python builtin max(seq_lens) and .item() on the result; iterating a tensor hands out one 0-d tensor per row, about 2 us each (0.35 ms at 128 rows, on the plan thread every decode step). Iterating the numpy view costs 0.1 us a row and the numpy scalar has .item() too. Everything else FlashInfer does with the row (.cpu(), .min(), len, a copy_ into its device buffer) is plain tensor behaviour.

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

Bases: object

Output of KVManager.plan for a single label

Parameters:

• cpu_indptrs (PagedIndptrs)

• views (list[SequenceView] | None)

• cuda_indptrs (PagedIndptrs | None)

cpu_indptrs

cuda_indptrs

decode_offsets

decode_pages

get_total_len()

property is_decode: bool

rows

property 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, kv_lens=None)

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)

• kv_lens (Tensor | None)

kv_lens: Tensor | None

Alias for field number 4

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.PinnedIndexRing(numel=4096, depth=2, per_plan=4)

Bases: object

A ring of page-locked int32 host buffers for the index arrays a plan hands FlashInfer, so its copy_(non_blocking=True) calls are really asynchronous (from pageable memory the driver stages each one).

The ring holds depth buffer sets, one per plan. begin_plan moves to the next set, and each take until the next begin_plan fills a buffer of that set no other take uses (the set grows if needed), so every label of a plan keeps its arrays (four takes each) however many labels there are. A set is rewritten only after the copies out of it have completed: the next begin_plan records an event on the stream the plan ran on (by then the plan’s copies, and its consumers’, were issued), and a set waits on its event before reuse. Plans do not overlap, so one pending set suffices. Without CUDA the buffers are plain tensors and nothing is recorded.

Parameters:

• numel (int)

• depth (int)

• per_plan (int)

begin_plan()

Return type:

None

take(values)

Return type:

Tensor

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, ring=None)

Parameters:

• segments (list[SequenceView])

• page_size (int)

• ring (PinnedIndexRing | None)

Return type:

PagedIndptrs

mstar.engine.resources.kv.plan.context_only_views(views, query_spans, page_size)

The views of a step’s streams without the tokens the step appends: length shrinks by to_compute (the pages beyond it dropped) and to_compute becomes the row’s query count. For attention over the stored context alone (AttentionStep.context_only).

Parameters:

• views (list[SequenceView])

• query_spans (list[int])

• page_size (int)

Return type:

list[SequenceView]

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