[diffusion] In-place LoRA merge/unmerge under layerwise offload

[niehen6174]

Motivation

Memory limits on MiniMax-H3

• GPU. The BF16 DiT does not fit on a 24 GB card. Serving uses layerwise offload: weights live on the host, and only the current layer is streamed onto the GPU.
• Host RAM. The offload store is already the full weight set (mmap / pin budget). Cloning a second unmerge snapshot is another 38 GB of anonymous memory on this DiT.
• The previous runtime LoRA path. set_lora, merge_lora_weights, unmerge_lora_weights, and deactivate_lora_weights all turned offload off and called load_all so they could see dense GPU tensors. That pulls the whole DiT onto the card and OOMs. Clearing or switching an adapter hits the same path: unmerge / deactivate restore W the same way set_lora writes it.

We cannot merge or unmerge by materializing the model, and we cannot afford a second host copy. Both writes have to run in place against the existing CPU store, one layer at a time.

Relation to #36062

#36062 already fixed the startup --lora-path case. In-place merge copy-on-writes the checkpoint mapping, and the LoRA wrapper used to clone a full base snapshot for unmerging. Together that peaked at 75.3 GiB of anonymous host memory on MiniMax-H3 + 4-step turbo LoRA. That PR copy-merges into a file-backed cache and keeps zero-copy unmerge views, so peak host anon falls to 19.2 GiB.

It explicitly left the runtime APIs on the in-place path. This PR is that remaining surface:

API	What it does
set_lora	load / apply / switch an adapter
merge_lora_weights	merge the currently loaded adapter into W
unmerge_lora_weights	restore the base (W back to pre-LoRA); required before switching
deactivate_lora_weights	clear LoRA: unmerge if needed, then disable the wrapper

HTTP: POST /v1/set_lora, /v1/merge_lora_weights, /v1/unmerge_lora_weights. All four go through _weight_update_context and the same CPU writeback. --lora-path stays on the #36062 cache.

---

Approach

Keep layerwise on for every runtime LoRA write (merge and unmerge / deactivate). For each layer, read the manager CPU view, run the update on GPU, write the result back with a blocking copy_, then release the GPU copy so the next prefetch reloads the new W:

src = CPU view (or live GPU weight if it is full-sized)
data = src.to(gpu)
data += scale * (B @ A)
dest.copy_(merged) into the CPU store   # blocking
release the stale GPU layer

Under offload the live layer.weight is a (1,) placeholder, so the write must hit the CPU store, not that placeholder. After wrap the live name becomes *.base_layer.weight while the manager still records *.weight; we alias that pair, rebind the manager to the same Parameter, and never replace the Parameter object. Unmerge snapshots stay zero-copy views of the CPU store (snapshot_base=False).

Applies to any BaseLayerWithLoRA + dense weight. Not GGUF/qweight, nested wraps, or DTensor.

---

Experiments

Dynamic set_lora / unmerge (this PR)

4090 D 24 GB, layerwise dit,text_encoder, MiniMax-H3 BF16 + lightx2v 4-step LoRA. set_lora after load (not --lora-path), then 5-step T2VA.

• Unit: aliases, park skips wrapped layers, writeback keeps the Parameter, view unmerge restores the base.
• E2E: set_lora does not load_all / OOM; the clip is a normal scene, not a color-noise grid.
• merge / unmerge / deactivate share the same writeback. Unmerge is unit-tested; no extra 5-step clip after unmerge.

---

Performance

4090 D 24 GB, 4-step LoRA via set_lora:

Stage	Time
set_lora (208 layers)	5.9 s
denoise	~99 s
generate	~122 s

A multi-stream H2D / compute / D2H pipeline was tried and dropped. Real merge is already a few seconds; overlapping transfers saved ~1 s and added a lot of code. Denoise still uses its existing prefetch copy_stream — that is inference, not LoRA merge.

---

CI States

Latest PR Test (Base): ❌ Run #33313683328
Latest PR Test (Extra): ❌ Run #33313683133
Latest PR Test (AMD ROCm 7.2): ❌ Run #33313683293