mstar.engine.resources.lora.backend

mstar.engine.resources.lora.backend#

Kernels that add a batch’s LoRA deltas: shrink h = x @ A^T, then expand y += s * h @ B^T.

Both read the step’s LoraAddressing from the pool, never a model-side index, so the eager and captured paths address slots the same way. Shapes, per layer of one target:

  • x [T, in], y [T, out] (written in place), h [T, slices * R]

  • a [slots, slices * R, in], b [slots, out, R], ranks / scalings [slots]

out_slices are this rank’s widths of a fused projection’s outputs ((q, k, v)); slice i reads h[:, i*R:(i+1)*R] and writes its own columns of y.

An embedding target’s shrink is a gather instead (gather_shrink): its input is a token id, in is the vocabulary, and x @ A^T of the id’s one-hot is A’s column for it.

Functions

expand_reference_(y, h, b, token_order, ...)

y[:, slice] += scaling * h[:, slice ranks] @ b[slot, slice]^T in place, as shrink_reference computes it.

gather_shrink_reference(ids, a, token_order, ...)

Plain torch for gather_shrink: fixed shapes and no host sync, so it compiles and captures on any device.

make_backend(kind, device)

The custom ops (ops.LoraOpsBackend) on every device, so torch.compile takes each kernel as one opaque node.

shrink_reference(x, a, token_order, ...)

h [T, slices * R] for x [T, in] against a [slots, slices * R, in], every token against every slot and then its own selected: no host sync and no data-dependent shape, so it compiles and captures, at slots times a dense LoRA's work.

Classes

LoraBackend()

class mstar.engine.resources.lora.backend.LoraBackend[source]#

Bases: ABC

apply_(x, y, a, b, addressing, ranks, scalings, out_slices, embedding=False)[source]#

Shrink (a gather for an embedding target, whose x is token ids) then expand into y; returns y.

Parameters:
Return type:

torch.Tensor

apply_linear_(x, y, a, b, addressing, ranks, scalings, out_slices)[source]#

Shrink then expand into y, for a linear target; a backend may fuse them.

Parameters:
Return type:

None

abstractmethod expand_fused(h, b, addressing, ranks, scalings, y, out_slices)[source]#

y[t, slice i] += scalings[slot(t)] * h[t, slice i] @ b[slot(t), slice i]^T.

Parameters:
Return type:

None

gather_shrink(ids, a, addressing, ranks)[source]#

h[t] = a[slot(t)][:, ids[t]], the shrink of a token’s one-hot; an id outside [0, in) (e.g. an image position’s -1) contributes nothing.

Parameters:
Return type:

torch.Tensor

abstractmethod shrink(x, a, addressing, ranks, num_slices=1)[source]#

h[t] = x[t] @ a[slot(t)]^T. Rows of rank-0 slots, and columns past a slot’s rank, may hold anything: the expand does not read them.

Parameters:
Return type:

torch.Tensor

mstar.engine.resources.lora.backend.expand_reference_(y, h, b, token_order, slot_indptr, ranks, scalings, out_slices)[source]#

y[:, slice] += scaling * h[:, slice ranks] @ b[slot, slice]^T in place, as shrink_reference computes it.

Parameters:
Return type:

None

mstar.engine.resources.lora.backend.gather_shrink_reference(ids, a, token_order, slot_indptr, ranks)[source]#

Plain torch for gather_shrink: fixed shapes and no host sync, so it compiles and captures on any device.

Parameters:
Return type:

Tensor

mstar.engine.resources.lora.backend.make_backend(kind, device)[source]#

The custom ops (ops.LoraOpsBackend) on every device, so torch.compile takes each kernel as one opaque node. "triton" ("auto" on CUDA) registers the Triton kernels for them; "torch" registers none. A deployment’s pools share one lora: block, so they choose alike, but registration is process-wide: the reference runs only while nothing else in the process has registered the Triton kernels.

Parameters:
Return type:

LoraBackend

mstar.engine.resources.lora.backend.shrink_reference(x, a, token_order, slot_indptr, ranks, num_slices)[source]#

h [T, slices * R] for x [T, in] against a [slots, slices * R, in], every token against every slot and then its own selected: no host sync and no data-dependent shape, so it compiles and captures, at slots times a dense LoRA’s work. Zeros where the expand does not read (dead tokens, columns past a slot’s rank).

Parameters:
Return type:

Tensor