mstar.engine.resources.lora.pool#
Per-request LoRA adapters resident in a fixed set of device slots.
The pool owns the slot weights, decides which adapter sits in which slot, and says per step
which slot each token addresses; a LoraBackend runs the kernels against that.
Slot 0 is the zero adapter (rank 0, scaling 0) and is never handed out: base-model rows and a
capture bucket’s padding rows address it, and kernels skip it. A request takes the slot of
its adapter on its first admit and holds it until it is removed, so a slot whose refcount is
zero is addressed by no live request. That is what makes it safe to overwrite: the worker
defers remove_request until no in-flight step references the request, so a refcount-0
slot is read by no step still on the GPU, and a slot taken by an admit that was rolled back
was never read at all.
Under tensor parallelism every rank of a lockstep instance keeps its own pool and makes
the same slot decisions with no journal: they depend only on the order of admits, which
the leader’s broadcast fixes, and on which teardowns landed between two admits, which the
worker fixes by landing each between the same two admits on every rank (after_tp_seq).
The one decision that can differ per rank, whether an adapter loaded, is agreed across the
instance before anything is installed.
A full pool is backpressure, not a refused step: a request on an adapter waits at the
runner’s admission (ask / reserve) until a slot can take its adapter, in one queue
with every other resource that asks, so it holds no KV pages meanwhile and the rows already
running never lose a step to it. A reservation counts slots and pins none: the slot an
admit picks follows from the slot table alone, and the reservations never promise more
slots than exist, so every reserved row finds one at its admit. A row whose adapter
failed to read and is in no slot is the exception: it is let through at once, and its
admit fails it before handing out any slot, so it takes none that a counted row needs.
Only the rank that chooses batches decides (a lockstep instance’s leader), so no other
rank keeps reservations; the admit refusal stays as a backstop.
A model with LoRA specs on several nodes (an LLM and its vision encoder) has a pool per node, and each admits only its own node’s rows. A request holds a slot in each until it is removed (none in a pool whose targets its adapter has nothing for), but where the graph runs the encoder once, before the LLM (Qwen3.5), it takes them in that order and a row holding an LLM slot never waits on the encoder’s pool, so the two pools’ waits cannot form a cycle.
Files are read on the pool’s reader thread, never on the GPU thread or under the lock: a request names an adapter, and every rank holds it in the host tier and starts reading it as it ingests the request; the leader reserves a slot for a request only once its own copy is in. A follower whose copy is still in flight when the leader’s step admits waits for it outside the lock, once the instance has agreed that no rank’s read failed. Its read usually started with the leader’s, at the same ingest, but not always: a tier keeps what the ingests and removes on its rank leave it, and nothing orders the two against each other across ranks (a follower takes removes from the leader, ingests from the conductor). A remove that races an ingest can leave a follower without a copy the leader kept, and the next request on that adapter then waits there for a whole read.
Adapters are also loaded and unloaded at runtime (load_adapter / unload_adapter),
which each rank applies whenever its control message arrives, not at a seq every rank
agrees on. That is safe because neither touches what an admit decides from (slot
refcounts, which slot holds an adapter, the LRU clock): a load only makes the adapter
servable (the host tier holds it from its first request), and an unload only forgets the
adapter’s name, its host copy and its name -> slot entry, once no live request names it.
An admit consults that entry only for a request naming the adapter, and the next such
request is ingested after a reload, which reaches every rank after the unload did. The
slot keeps the old weights, unheld, until LRU eviction reuses it.
Classes
|
|
|
Which slot each token of this step addresses, grouped by slot: the pool's plan output, and all a backend reads. |
|
|
|
- class mstar.engine.resources.lora.pool.LoraAdapterPool(spec, device, rank=0, world_size=1, backend=None, dtype=None, agree=None, decides=True)[source]#
Bases:
Resource- Parameters:
spec (LoraSpec)
device (torch.device | str)
rank (int)
world_size (int)
backend (LoraBackend | None)
dtype (torch.dtype | None)
decides (bool)
- addressing(grouping=1)[source]#
This step’s addressing for rows of
grouping(LoraTarget.row_grouping), for a backend to run against.- Parameters:
grouping (int)
- Return type:
- admit(step, ctx)[source]#
Give each row whose adapter is not resident yet a slot, loading it there.
Nothing to do once every row holds its slot, which is every step after a request’s first, so the steady state costs one set check.
Rows whose adapter is resident take its slot before any miss picks a victim, so a miss cannot evict an unheld adapter that a later row of the same step names.
Every miss’s slot is picked in that same locked section, before any wait: what is picked follows only from the slot table at the admit’s start, which is the same on every rank of a lockstep instance, whatever lands while a rank waits for its copy or for the load agreement. The ranks then agree, in one collective, where each miss’s copy stands, before any rank waits for one. Rows whose adapter failed to read on any rank fail before any slot is handed out, so a row let through on a failed read never takes the slot a counted row was promised.
- Parameters:
step (LoraStep)
ctx (StepContext)
- Return type:
- ask(rid, node_name, graph_walk)[source]#
ASKS for a row that will need a slot, or the outcome that fails a row whose adapter failed to read. None on a rank that doesn’t decide, and for a base row, a row that holds its slot or a reservation (asked again by another node the pool serves), and an adapter with nothing for this pool.
- Return type:
AdmitOutcome | None
- base_only()[source]#
A forward outside any step (a compile’s warm-up) on this thread, as a base-only step: its deltas add nothing. A compile traced here takes the base-only branch, so the first step with an adapter row compiles its own.
- Return type:
Iterator[None]
- classmethod build(spec, info)[source]#
- Parameters:
spec (LoraSpec)
info (EngineResourceInfo)
- Return type:
- commit(step, ctx)[source]#
record step consumption
- Parameters:
step (LoraStep)
ctx (StepContext)
- Return type:
None
- 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
planholds one per wrapper, which is why the attention managers key theirs per slot. The main runner already double-buffers whenever a resourcesupports_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_requestfence 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]#
Open the request’s state. The worker calls this once per request, on its first NewRequest, though one comes per partition it serves.
- Parameters:
overrides (ResourceReqConfig | None)
- layer_delta(target, layer_idx)[source]#
What a layer calls to add
target’s delta at stack layerlayer_idx.- Parameters:
- Return type:
- layer_ranks(target, layer_idx)[source]#
[slots]: each slot’s rank at one target and stack layer, 0 where its adapter adds nothing there.
- load_adapter_async(name, path, digest=None)[source]#
Serve
namefrompath: read it on the reader thread, check it against the spec (and againstdigest, the weights the API server checked), and only then make it servable. The read is not kept, since the host tier holds only what requests name: the first request on it reads it again as it is ingested. The future fails if it is served already, is still retiring from an unload, or cannot be read; a refusal visible now raises here. A request that names it takes a slot on its first admit, as for an adapter configured at startup.
- plan(step, ctx)[source]#
ret is immutable and opaque to runner; only gives to ctx.plan_results
- Parameters:
step (LoraStep)
ctx (StepContext)
- Return type:
LoraAddressing | None
- reserve(rid, node_name, graph_walk, published)[source]#
Count
rid’s adapter against the free slots once its host copy is in, so its admit does not wait for a read on this rank. A row whose adapter failed to read and is in no slot is let through at once, slots or not: its admit fails it, on every rank alike. One whose adapter is still resident may yet run on that slot, so it counts like any other.- Return type:
int | None
- rollback_admit(step, ctx)[source]#
Give back the slots this step’s admit took. An adapter it loaded stays resident in its now-unheld slot, as a cache.
- Parameters:
step (LoraStep)
ctx (StepContext)
- Return type:
None
- runs_adapters(rids)[source]#
Whether any of
ridsnames an adapter with something for this pool: set at ingest and from the adapters’ files, so the same on every rank of a lockstep instance.- Return type:
- property store: LoraAdapterStore#
- unload_adapter(name)[source]#
Stop serving
nameto new requests. Its host copy and its name -> slot entry are forgotten now if no live request names it, else when the last one is removed; those requests run to the end on its weights. Returns whether it was forgotten now. An adapter this pool does not serve is a no-op (a rolled-back load that failed here).
- class mstar.engine.resources.lora.pool.LoraAddressing(token_order, slot_indptr)[source]#
Bases:
objectWhich slot each token of this step addresses, grouped by slot: the pool’s plan output, and all a backend reads.
Views of one fixed device buffer per capture slot (or of a reused eager buffer), sized to this step’s tokens, which are packed in segment order. A token past every slot’s range (a capture bucket’s tail) addresses none.
A target that sees other rows than the step’s tokens (one per segment, or one per
token_foldtokens;LoraTarget.row_grouping) reads an addressing of its own, whose “tokens” are those rows.
- class mstar.engine.resources.lora.pool.LoraLayerDelta(pool, target, layer_idx)[source]#
Bases:
objecty += delta(x)for one target linear at one stack layer, over the current step.A layer computes its base output, then calls this before anything consumes
y— for a row-parallel linear, before its all-reduce. Rows on the base model address the zero adapter, which the kernels skip; a step with no adapter row at all returnsyuntouched (LoraAdapterPool.active). The layer’s slot weights are views resolved once, here; a call reads only the step’s addressing.- Parameters:
pool (LoraAdapterPool)
target (str)
layer_idx (int)