mstar.engine.resources.position.manager

mstar.engine.resources.position.manager#

does plan_rope/apply_rope of cache_manager.py

has ownership of per-(request,label) position counter

Classes

PositionManager()

PublishedPositionInfo([counters])

Where each of a request's streams has reached.

RopeManager(config, device[, max_seq_len, ...])

rotary position on top flashinfer in-place kernels
class mstar.engine.resources.position.manager.PositionManager[source]#

Bases: Resource

classmethod build(spec, info)[source]#
Parameters:
prefix_skip_safe: bool = True#
class mstar.engine.resources.position.manager.PublishedPositionInfo(counters=<factory>)[source]#

Bases: PublishedInfo

Where each of a request’s streams has reached. label -> next position.

Parameters:

counters (dict[str, int])

clone()[source]#

Return independent publication metadata for deferred sending.

Return type:

PublishedPositionInfo

counters: dict[str, int]#
update(other)[source]#
Parameters:

other (PublishedPositionInfo)

Return type:

None

class mstar.engine.resources.position.manager.RopeManager(config, device, max_seq_len=32768, head_dim=None, dtype=torch.bfloat16)[source]#

Bases: PositionManager

rotary position on top flashinfer in-place kernels

Parameters:
admit_retrieve(rid, node_name, graph_walk, published)[source]#

Take on the counters a stream was published with.

A node that receives another node’s KV (a CFG branch on its own GPU, a decode engine after disaggregated prefill) has never advanced these streams itself, so its counters start at zero while the KV it just pulled in is many tokens long. Monotonic, like the KV lengths beside them: a stale echo of this node’s own publish never rewinds it.

Parameters:
Return type:

AdmitOutcome

apply_cached_prefix(rid, node_name, graph_walk, inputs, prefix)[source]#

Move the stream’s counter past what the cache holds of this walk.

The counter moves here and not at admit. A walk the cache serves never admits, and the walk after it can place its positions in prepare_inputs, before its own admit, so a counter seeded at admit would put that walk on top of the served ones. The position is absolute, so a repeated probe moves nothing.

Parameters:
Return type:

None

apply_qk(q, k, label, rotary_dim=None, interleave=None, rope_scale=None, rope_theta=None, rope_dtype=None, **kwargs)[source]#
Parameters:
clear_preplan()[source]#
commit(step, ctx)[source]#

record step consumption

Parameters:
depends_on()[source]#
fingerprint()[source]#

What this resource contributes to the prefix cache’s root.

Everything a stored page depends on beyond its own tokens. None leaves the resource out of the root, which is right only where it cannot change what a page holds.

Return type:

bytes

property force_double_buffer#

Whether this resource needs double-buffering even off the pre-plan path.

A resource that stages a step’s layout into a reused host buffer and issues a non-blocking H2D into a graph-read device buffer has a race the moment the CPU runs ahead of the GPU: plan(N+1) can overwrite the buffer before step N’s DMA has retired, and the replay attends with N+1’s data. The buffer need not be ours — FlashInfer’s plan holds one per wrapper, which is why the attention managers key theirs per slot. The main runner already double-buffers whenever a resource supports_preplan; this flag extends that to a runner that has no pre-plan path but still replays such a resource — notably the piecewise runner — and makes _exec_per_request fence before it reuses a slot, since its sub-steps all run inside one step. Off by default; a resource opts in only if it has this hazard.

ingest_request(rid, overrides=None)[source]#

Must be idempotent: the conductor sends one NewRequest per partition, so a worker serving several partitions of a request ingests it once for each, possibly after a node has already run a step for it.

Parameters:

rid (str)

plan(step, ctx)[source]#

submits one position id vector for each plan label

plan labels and token ordering comes from KV plan output

Parameters:
Return type:

dict[str, Tensor]

pos_ids(label)[source]#
Parameters:

label (str)

Return type:

Tensor | None

position(rid, label)[source]#
Parameters:
Return type:

int

publish(request_id)[source]#
Parameters:

request_id (str)

Return type:

PublishedPositionInfo | None

publish_from_snapshot(request_id, snapshot, node_name, graph_walk)[source]#

Finish a publish_snapshot_for_step record into the published info.

Parameters:

  • request_id (str)

  • node_name (str | None)

  • graph_walk (str | None)

publish_snapshot_for_step(request_id, node_name, graph_walk)[source]#

A cheap record of what publish_for_step would export right now.

The worker finishes it with publish_from_snapshot only for the requests whose frame or completion carries the publication, which in a steady decode loop is none of them. The default takes the eager publication itself, so a resource is correct without overriding; one whose publish builds something per request (a page list, a transfer descriptor) overrides both to record a few scalars here instead.

Parameters:

  • request_id (str)

  • node_name (str | None)

  • graph_walk (str | None)

remove_request(rid)[source]#
Parameters:

rid (str)

reset_request(rid, free=False)[source]#

For clearing dummy RIDs during cuda graph capture

Parameters:
property supports_preplan#