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)

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

Bases: Resource

classmethod build(spec, info)[source]#
Parameters:
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])

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

other (PublishedPositionInfo)

Return type:

None

class mstar.engine.resources.position.manager.RopeManager(config, device)[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_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]#
ingest_request(rid, overrides=None)[source]#
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

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#