Chapter 01 — Shaving Milliseconds Off LoRA Inference for Diffusion Models

Most write-ups on LoRA stop at training: rank r, target modules, alpha scaling, done. The part nobody warns you about is what happens the moment you want to swap adapters at inference time — a user picks a different style, or your serving layer wants to batch requests that each need a different fine-tune, and suddenly the naive approach falls apart.

The naive approach, and why it's slow

The obvious thing to do is merge the LoRA weights into the base weights before every request:

def merge_lora(base_weight, lora_a, lora_b, alpha, rank):
    # base_weight: [out, in], lora_a: [rank, in], lora_b: [out, rank]
    delta = (lora_b @ lora_a) * (alpha / rank)
    return base_weight + delta

This is fine if you only ever serve one adapter per process. It falls apart the moment adapters are a per-request concern: merging is an out x in matrix add on every linear layer you've adapted, and for a UNet with a few hundred attention/conv layers, re-merging on every single request adds up to real, measurable latency — I was seeing 40-60ms of pure merge overhead before the actual denoising loop even started, on a workload where the diffusion step itself was budgeted at ~800ms for 20 steps.

What actually worked: keep the delta separate

Instead of merging into the base weight, keep the low-rank delta as its own pair of small matmuls and add the result at the activation level:

class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, rank: int, alpha: float):
        super().__init__()
        self.base = base
        self.lora_a = nn.Parameter(torch.zeros(rank, base.in_features))
        self.lora_b = nn.Parameter(torch.zeros(base.out_features, rank))
        self.scale = alpha / rank
        self.base.weight.requires_grad_(False)
 
    def forward(self, x):
        out = self.base(x)
        # rank is small (4-64), so this is cheap relative to the base matmul
        out = out + (x @ self.lora_a.T @ self.lora_b.T) * self.scale
        return out

This looks like it should be slower — two extra matmuls per adapted layer instead of zero — but the rank is tiny (I was using r=16 almost everywhere), so x @ lora_a.T collapses the feature dimension down to 16 before the second matmul blows it back up. The FLOPs are negligible next to the base linear layer, and critically, you never touch the base weight tensor, so:

• Adapter swaps are just pointer swaps on a small pair of tensors, not a full-model weight mutation.
• You can batch requests using different adapters in the same forward pass by keeping a per-sample adapter index and gathering the right lora_a/ lora_b pair per row, instead of one adapter per process.
• Base weights stay untouched in VRAM, so multiple adapters can share a single resident copy of the UNet.

Batching across adapters

The part that actually moved the needle in production was batching requests that use different adapters together, rather than routing each adapter to its own queue:

def batched_lora_forward(x, base_linear, adapters, adapter_ids):
    # adapters: dict[str, (lora_a, lora_b, scale)]
    # adapter_ids: LongTensor[batch] mapping each row to an adapter
    out = base_linear(x)
    for name, (lora_a, lora_b, scale) in adapters.items():
        mask = adapter_ids == name_to_id[name]
        if not mask.any():
            continue
        delta = (x[mask] @ lora_a.T @ lora_b.T) * scale
        out[mask] += delta
    return out

It's not elegant — a scatter/gather over a Python dict loop is not going to win any style points — but it turns "one adapter per GPU worker" into "one GPU worker serving N adapters concurrently," which is the difference between needing N GPUs and needing one. For a catalog of ~30 style adapters with uneven traffic, that's the whole ballgame.

The lesson that generalized past this one project: the fine-tuning technique and the serving strategy are separate design problems. LoRA's low-rank structure is what makes cheap, mid-request adapter swaps possible at all — but you only get that benefit if your serving code is written to keep the delta separate, batch across adapters, and never assume "one model weight tensor per process."