mstar.engine.resources.lora.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.

mstar::lora_expand_ adds into y in place and says so (mutates_args): functionalization then keeps every read of y on the right side of the add and never drops it, whatever Inductor reorders around it.

The kernel registered here is the plain-torch reference (backend.shrink_reference / expand_reference_), on any device, with no host sync and no data-dependent shape; triton_backend registers the CUDA kernels over it. Both read the same arguments with the same meaning, so the fake (shape) implementations and torch.library.opcheck can be checked on a CPU.

Rows of h the expand does not read (tokens on the zero adapter, columns past a slot’s rank) are unspecified; the reference writes zeros there.

Classes

LoraOpsBackend()

The custom ops: on CUDA the Triton kernels once triton_backend is imported, elsewhere the reference above, so a CPU test runs exactly the graph a GPU compiles.

class mstar.engine.resources.lora.ops.LoraOpsBackend

Bases: LoraBackend

The custom ops: on CUDA the Triton kernels once triton_backend is imported, elsewhere the reference above, so a CPU test runs exactly the graph a GPU compiles.

apply_linear_(x, y, a, b, addressing, ranks, scalings, out_slices)

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

Parameters:

addressing (LoraAddressing)

expand_fused(h, b, addressing, ranks, scalings, y, out_slices)

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

Parameters:

addressing (LoraAddressing)

gather_shrink(ids, a, addressing, ranks)

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:

addressing (LoraAddressing)

shrink(x, a, addressing, ranks, num_slices=1)

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:

addressing (LoraAddressing)