vllm_omni.diffusion.lora.layers.base_linear ¶

DiffusionBaseLinearLayerWithLoRA ¶

Bases: BaseLinearLayerWithLoRA

Diffusion-specific base that overrides apply() to use direct torch matmul instead of punica_wrapper.

punica_wrapper is used to hold multiple LoRA slots and slices efficiently.

This matches the semantics of PunicaWrapperGPU.add_lora_linear(): - Shrink: buffer = (x @ lora_a.T) - Expand: y += buffer @ lora_b.T

All other functionality (weight management, TP slicing, forward logic) is inherited from vLLM's BaseLinearLayerWithLoRA.

lora_a_stacked `instance-attribute` ¶

lora_a_stacked: tuple[Tensor, ...]

lora_b_stacked `instance-attribute` ¶

lora_b_stacked: tuple[Tensor, ...]

apply ¶

apply(x: Tensor, bias: Tensor | None = None) -> Tensor

override: Use simple matmul instead of punica_wrapper.add_lora_linear().

This matches the exact computation in PunicaWrapperGPU.add_lora_linear() for the single-LoRA case. For packed projections (e.g. fused QKV), we apply LoRA per-slice using output_slices.

create_lora_weights ¶

create_lora_weights(
    max_loras: int, lora_config, model_config=None
) -> None

reset_lora ¶

reset_lora(index: int)

resume_lora ¶

resume_lora() -> None

Restore the mask captured by suspend_lora().

A layer registered after the suspend, or one whose buffers were torn down meanwhile, has no mask to restore and keeps its own inactive one.

set_lora ¶

set_lora(
    index: int,
    lora_a: Tensor | list[Tensor | None],
    lora_b: Tensor | list[Tensor | None],
)

suspend_lora ¶

suspend_lora() -> None

Gate every slice off without dropping the uploaded weights.

apply() reads the stacked buffers only for slices this mask marks active, so a suspended adapter costs nothing while staying resident. A saved mask means the layer is already suspended; keep the original rather than overwriting it with the all-False mask now in place.