mstar.engine.resources.lora

class mstar.engine.resources.lora.AdapterInfo(rank, unmapped, misfits=())

Bases: NamedTuple

What validation needs to know about an adapter checkpoint, read without its weights.

rank covers every LoRA module in the file, those left to another spec included. misfits names modules whose shapes do not fit their target (an adapter trained on another size of the model).

Parameters:

• rank (int)

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

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

misfits: tuple[str, ...]

Alias for field number 2

rank: int

Alias for field number 0

unmapped: tuple[str, ...]

Alias for field number 1

class mstar.engine.resources.lora.LoraBaseOnly(key=None)

Bases: object

The capture key of a step none of whose rows names an adapter: the node’s own cg_key_info (key), wrapped. A node serving adapters captures each bucket under both keys, and the base-only graph records no LoRA kernel: at decode width the launches of a delta that adds zero cost about what the work does.

Parameters:

key (Any)

key: Any = None

class mstar.engine.resources.lora.LoraModuleRef(target, layer, shard_id=None)

Bases: NamedTuple

Where one checkpoint module’s adapter lands: a target, a stack layer, and the stacked-shard id of its slice (None for a single-slice target).

Parameters:

• target (str)

• layer (int)

• shard_id (str | int | None)

layer: int

Alias for field number 1

shard_id: str | int | None

Alias for field number 2

target: str

Alias for field number 0

class mstar.engine.resources.lora.LoraReqConfig(adapter=None)

Bases: ResourceReqConfig

The adapter a request runs under; None runs the base model.

Parameters:

adapter (str | None)

adapter: str | None = None

validate(spec=None)

Raise ValueError if a client-supplied value is out of range, or one the resource built from spec cannot honour.

The conductor calls it at admission, so a bad request is a 400 before any worker sees it.

Parameters:

spec (LoraSpec | None)

Return type:

None

class mstar.engine.resources.lora.LoraSpec(resource_key, nodes, num_layers, targets, module_map, max_slots=4, max_rank=16, cpu_capacity=16, adapters=<factory>, adapter_digests=<factory>, dtype=None, backend='auto', convert_keys=None, ignored_modules=None, dropped_targets=<factory>)

Bases: NodeResourceSpec

Per-request LoRA adapters for the linears named in targets.

max_slots adapters are resident on the device at once, beside a zero adapter that base-model rows use, and up to cpu_capacity are kept pinned on the host. adapters maps the names requests ask for to safetensors paths. module_map places each checkpoint module on a target; a node’s spec comes from module_lora_spec (mstar.model.components.lora.modules), whose ModuleLoraMap is what binds the node’s modules (a spec with any other resolver is refused as the deployment’s specs are declared, declared_lora_specs). convert_keys rewrites a model-native checkpoint layout first.

Parameters:

• resource_key (str)

• nodes (set[str])

• num_layers (int)

• targets (dict[str, LoraTarget])

• module_map (Callable[[str], LoraModuleRef | None])

• max_slots (int)

• max_rank (int)

• cpu_capacity (int)

• adapters (dict[str, str])

• adapter_digests (dict[str, bytes])

• dtype (dtype | None)

• backend (Literal['auto', 'torch', 'triton'])

• convert_keys (Callable | None)

• ignored_modules (Callable[[str], bool] | None)

• dropped_targets (frozenset[str])

adapter_digests: dict[str, bytes]

adapter_info(name)

Rank and unmapped modules of a configured adapter, read once and cached.

Parameters:

name (str)

Return type:

AdapterInfo

adapters: dict[str, str]

add_adapter(name, path, digest=None)

Serve name from path, whose weights digest to digest, from now on (a runtime load).

Parameters:

• name (str)

• path (str)

• digest (bytes | None)

Return type:

None

apply_yaml_overrides(max_slots=None, max_rank=None, cpu_capacity=None, adapters=None, adapter_digests=None, targets=None)

How many adapters stay resident, the largest rank served, and which adapters this deployment serves. targets keeps only those of the model’s targets (one list for every LoRA spec of the model, so a name another spec declares is not an error here; serving_resources checks the list as a whole): dropping the LM head and embedding saves their vocabulary-sized slots, and an adapter that touches a dropped target is refused.

Every value is checked before any is set, so a refused override leaves the spec as declared.

Parameters:

• max_slots (int | None)

• max_rank (int | None)

• cpu_capacity (int | None)

• adapters (dict[str, str] | None)

• adapter_digests (dict[str, str] | None)

• targets (list[str] | None)

backend: Literal['auto', 'torch', 'triton'] = 'auto'

check_adapter(name, info)

Refuse an adapter this deployment cannot serve: no LoRA layers, a rank above max_rank, or a module outside the declared targets.

Parameters:

• name (str)

• info (AdapterInfo)

Return type:

None

convert_keys: Callable | None = None

cpu_capacity: int = 16

dropped_targets: frozenset[str]

dtype: dtype | None = None

ignored_modules: Callable[[str], bool] | None = None

max_rank: int = 16

max_slots: int = 4

module_map: Callable[[str], LoraModuleRef | None]

num_layers: int

property num_slots: int

Device slots, the zero adapter’s included.

remove_adapter(name)

Stop serving name (a runtime unload).

Parameters:

name (str)

Return type:

None

property resource_class: type[Resource]

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

stop_ignoring(modules)

Refuse, rather than ignore, the checkpoint modules modules matches: those of another node’s spec the deployment does not build, which no spec loads now.

Parameters:

modules (Callable[[str], bool])

Return type:

None

targets: dict[str, LoraTarget]

class mstar.engine.resources.lora.LoraStep(segments=None)

Bases: ResourceStep

One step’s rows: a segment per request, its span the tokens that request puts through the LoRA’d linears this step, in the order they are packed.

Parameters:

segments (tuple[Segment, ...] | None)

class mstar.engine.resources.lora.LoraTarget(in_features, out_slices, parallel='column', slice_names=(), slice_units=None, slice_packs=None, slice_align=None, layers=None, kind='linear', per_segment=False, token_fold=1)

Bases: object

One family of linears an adapter may touch, e.g. every layer’s fused qkv_proj.

Shapes are the unsharded linear’s. out_slices lists the output widths of a fused projection in order ((q, k, v)); each slice has its own A and B, named in a checkpoint by the stacked-shard id in slice_names. A single linear has one slice.

parallel says how the base linear is split across tensor-parallel ranks:

• "column": output rows are split, so A is replicated and B is sharded by output rows of each slice. The delta lands on this rank’s rows of the output.

• "row": input columns are split, so A is sharded along the input and B replicated. x @ A_r^T is then a partial sum over this rank’s input columns, and so is the expanded delta: add it to the partial base output before the layer’s all-reduce, which sums both. LoRA adds no collective.

• "none": every rank holds the whole adapter and adds the whole delta to an output that is already whole (after the base layer’s own collective).

slice_units counts the shardable units (heads) of each slice. A slice with fewer units than ranks is replicated across groups of ranks rather than split, as QKVParallelLinear does with GQA’s K/V heads. None treats every row as a unit.

slice_packs names checkpoint modules that pack several slices into one tensor: a gated delta-net’s in_proj_qkv is [q | k | v], three outputs the base shards across ranks on their own. Such a module has one A, which each of its slices takes, and a B whose rows are those slices’ in order.

slice_align rounds each slice’s rank-local width up to a multiple, for a base linear that pads its fused output per rank so the next slice starts aligned. The padding rows follow the slice’s own, and B is zero there.

layers names the stack layer indices this target exists on (a hybrid stack has attention on some layers only); None means all of the spec’s num_layers.

kind="embedding" is a token embedding: in_features is the vocabulary, a row’s input is its token id, and the shrink gathers A’s column for it (x @ A^T of the token’s one-hot), as PEFT’s lora_embedding_A / _B define it.

Which rows the linear sees, for the per-row adapter choice: one per token of the step (the default), one per segment (per_segment, e.g. an LM head over each row’s last token), or one per token_fold consecutive tokens of a segment (a patch merger that folds 2x2 patches into one row).

Parameters:

• in_features (int)

• out_slices (tuple[int, ...])

• parallel (Literal['column', 'row', 'none'])

• slice_names (tuple[str | int, ...])

• slice_units (tuple[int, ...] | None)

• slice_packs (tuple[tuple[str | int, tuple[str | int, ...]], ...] | None)

• slice_align (tuple[int, ...] | None)

• layers (tuple[int, ...] | None)

• kind (Literal['linear', 'embedding'])

• per_segment (bool)

• token_fold (int)

in_features: int

input_cols(rank, world_size)

(start, size) of rank’s input columns.

Parameters:

• rank (int)

• world_size (int)

Return type:

tuple[int, int]

kind: Literal['linear', 'embedding'] = 'linear'

layers: tuple[int, ...] | None = None

local_width(slice_idx, world_size)

Rows of slice slice_idx in one rank’s base output, padding included.

Parameters:

• slice_idx (int)

• world_size (int)

Return type:

int

module_slices(shard_id)

The slices a checkpoint module lands on, in the order its B’s rows hold them: several for a slice_packs name, else the one slice_index gives.

Parameters:

shard_id (str | int | None)

Return type:

tuple[int, …]

property num_slices: int

property out_features: int

out_slices: tuple[int, ...]

parallel: Literal['column', 'row', 'none'] = 'column'

per_segment: bool = False

property row_grouping: int

1 per token, k per k folded tokens, 0 per segment.

Type:

Which addressing the linear reads

shard(world_size)

This target as one rank of world_size sees it: its slices the widths they have in that rank’s base output, alignment padding included.

Parameters:

world_size (int)

Return type:

LoraTarget

shard_pair(slice_idx, a, b, rank, world_size)

rank’s share of one slice’s a [r, in] and b [out_slice, r].

Parameters:

• slice_idx (int)

• a (Tensor)

• b (Tensor)

• rank (int)

• world_size (int)

Return type:

tuple[Tensor, Tensor]

slice_align: tuple[int, ...] | None = None

slice_index(shard_id)

Which slice a checkpoint module lands on, from its stacked-shard id.

Parameters:

shard_id (str | int | None)

Return type:

int

slice_names: tuple[str | int, ...] = ()

slice_packs: tuple[tuple[str | int, tuple[str | int, ...]], ...] | None = None

slice_rows(slice_idx, rank, world_size)

(start, size) of rank’s output rows within slice slice_idx.

Parameters:

• slice_idx (int)

• rank (int)

• world_size (int)

Return type:

tuple[int, int]

slice_units: tuple[int, ...] | None = None

token_fold: int = 1

mstar.engine.resources.lora.lora_enabled(model_config)

Whether a deployment’s YAML turns LoRA on (a resources: lora: block). A resources: value that is not a mapping turns nothing on: the conductor’s validate_deployment reports it with the deployment’s other problems.

Parameters:

model_config (Mapping[str, Any])

Return type:

bool

mstar.engine.resources.lora.node_capture_key(key)

The node’s own cg_key_info in a capture key, without its LoraBaseOnly.

Parameters:

key (Any)

Return type:

Any

mstar.engine.resources.lora.resolve_adapter_path(path)

The safetensors file an adapter path names: the path itself, or a PEFT directory’s adapter_model.safetensors (what vLLM’s lora_path names).

Parameters:

path (str)

Return type:

str

mstar.engine.resources.lora.served_lora_adapters(model_config)

The adapter names a deployment’s YAML serves; empty when it does not turn LoRA on.

Parameters:

model_config (Mapping[str, Any])

Return type:

tuple[str, …]

Modules

backend	Kernels that add a batch's LoRA deltas: shrink h = x @ A^T, then expand y += s * h @ B^T.
checkpoint	Reading LoRA checkpoints: shared by the static merge and the per-request adapter pool.
config	What a model declares about per-request LoRA adapters.
identity	Which weights an adapter name stands for, as the prefix cache keys them.
ops	The LoRA shrink and expand as custom ops, so torch.compile sees one opaque node each, and lora_apply_, the two as one op, which is what a linear target runs.
pool	Per-request LoRA adapters resident in a fixed set of device slots.
serving	Which LoRA specs a deployment builds from its serving config's resources: lora: block: the model's declared specs, sized by the block, narrowed to the targets its adapters touch, and those left with no targets dropped.
store	The host tier of the adapter pool: adapters read, sharded for this rank, and pinned.