From 4927fd5412f22098a0f8aa8aa455a538d53e433c Mon Sep 17 00:00:00 2001 From: 0z5a Date: Thu, 3 Sep 2026 10:37:35 +0800 Subject: [PATCH] [Core] Optimize diffusion LoRA expand accumulation (#6268) Signed-off-by: 0z5a <0z5a@users.noreply.github.com> Signed-off-by: 0z5a Co-authored-by: 0z5a <0z5a@users.noreply.github.com> --- .../benchmark_diffusion_lora_expand.py | 156 ++++++++++++++++++ tests/diffusion/lora/test_base_linear.py | 74 ++++++++- .../diffusion/lora/layers/base_linear.py | 13 +- 3 files changed, 239 insertions(+), 4 deletions(-) create mode 100644 benchmarks/kernels/benchmark_diffusion_lora_expand.py diff --git a/benchmarks/kernels/benchmark_diffusion_lora_expand.py b/benchmarks/kernels/benchmark_diffusion_lora_expand.py new file mode 100644 index 000000000..0f44c0ce4 --- /dev/null +++ b/benchmarks/kernels/benchmark_diffusion_lora_expand.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Benchmark diffusion LoRA expand-and-accumulate implementations. + +This compares the previous functional expand path with direct accumulation into +the selected output slice. It is a synthetic kernel microbenchmark and does not +load a diffusion model. + +Example: + python benchmarks/kernels/benchmark_diffusion_lora_expand.py \ + --m 1,16,128,512,4096,16384 \ + --in-dim 4096 --out-dim 4096 --rank 64 \ + --dtype bfloat16 --packed --warmup 30 \ + --small-iterations 500 --large-iterations 100 +""" + +import argparse +import math +import statistics +from collections.abc import Callable + +import torch + +MiB = 1024**2 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--m", default="1,16,128,512,4096,16384") + parser.add_argument("--in-dim", type=int, default=4096) + parser.add_argument("--out-dim", type=int, default=4096) + parser.add_argument("--rank", type=int, default=64) + parser.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16") + parser.add_argument( + "--packed", + action="store_true", + help="Accumulate into the second half of a packed output tensor.", + ) + parser.add_argument("--warmup", type=int, default=30) + parser.add_argument("--small-iterations", type=int, default=500) + parser.add_argument("--large-iterations", type=int, default=100) + parser.add_argument("--large-m-threshold", type=int, default=1024) + return parser.parse_args() + + +def percentile(samples: list[float], probability: float) -> float: + ordered = sorted(samples) + position = probability * (len(ordered) - 1) + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + weight = position - lower + return ordered[lower] * (1 - weight) + ordered[upper] * weight + + +def measure_latency_ms(operation: Callable[[], None], warmup: int, iterations: int) -> list[float]: + for _ in range(warmup): + operation() + torch.accelerator.synchronize() + + starts = [torch.cuda.Event(enable_timing=True) for _ in range(iterations)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(iterations)] + for start, end in zip(starts, ends, strict=True): + start.record() + operation() + end.record() + torch.accelerator.synchronize() + return [start.elapsed_time(end) for start, end in zip(starts, ends, strict=True)] + + +def measure_extra_peak_mib(operation: Callable[[], None]) -> float: + torch.accelerator.synchronize() + torch.accelerator.reset_peak_memory_stats() + allocated = torch.accelerator.memory_allocated() + operation() + torch.accelerator.synchronize() + peak = torch.accelerator.max_memory_allocated() + return (peak - allocated) / MiB + + +def print_header(args: argparse.Namespace, dtype: torch.dtype) -> None: + device_name = torch.cuda.get_device_name() + print(f"device={device_name} torch={torch.__version__} dtype={dtype}") + print(f"in_dim={args.in_dim} out_dim={args.out_dim} rank={args.rank} packed={args.packed} warmup={args.warmup}") + print("M iters old_p10 old_p50 old_p90 new_p10 new_p50 new_p90 speedup old_MiB new_MiB max_abs") + + +def benchmark_shape(m: int, args: argparse.Namespace, dtype: torch.dtype) -> None: + device = torch.device("cuda") + output_width = args.out_dim * (2 if args.packed else 1) + offset = args.out_dim if args.packed else 0 + output_slice = slice(offset, offset + args.out_dim) + scale = 1 / math.sqrt(args.in_dim) + + x = torch.randn(m, args.in_dim, device=device, dtype=dtype) + lora_a = torch.randn(args.rank, args.in_dim, device=device, dtype=dtype) * scale + lora_b = torch.randn(args.out_dim, args.rank, device=device, dtype=dtype) / math.sqrt(args.rank) + initial_output = torch.randn(m, output_width, device=device, dtype=dtype) + + def baseline(output: torch.Tensor) -> None: + buffer = x @ lora_a.t() + delta = buffer @ lora_b.t() + selected = output[:, output_slice] + selected[:] = selected + delta + + def optimized(output: torch.Tensor) -> None: + buffer = x @ lora_a.t() + selected = output[:, output_slice] + torch.addmm(selected, buffer, lora_b.t(), out=selected) + + reference = initial_output.clone() + candidate = initial_output.clone() + baseline(reference) + optimized(candidate) + torch.testing.assert_close(candidate, reference, rtol=1e-2, atol=1e-2) + max_abs = (candidate - reference).abs().max().item() + + old_output = initial_output.clone() + new_output = initial_output.clone() + iterations = args.large_iterations if m >= args.large_m_threshold else args.small_iterations + + old_peak = measure_extra_peak_mib(lambda: baseline(old_output)) + new_peak = measure_extra_peak_mib(lambda: optimized(new_output)) + old_samples = measure_latency_ms(lambda: baseline(old_output), args.warmup, iterations) + new_samples = measure_latency_ms(lambda: optimized(new_output), args.warmup, iterations) + + old_p50 = statistics.median(old_samples) + new_p50 = statistics.median(new_samples) + print( + f"{m:<7} {iterations:<6} " + f"{percentile(old_samples, 0.1):>8.5f} {old_p50:>8.5f} {percentile(old_samples, 0.9):>8.5f} " + f"{percentile(new_samples, 0.1):>8.5f} {new_p50:>8.5f} {percentile(new_samples, 0.9):>8.5f} " + f"{old_p50 / new_p50:>8.3f} {old_peak:>8.2f} {new_peak:>8.2f} {max_abs:>8.5f}" + ) + + +def main() -> None: + args = parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("This benchmark requires a CUDA GPU") + if args.warmup < 0 or args.small_iterations <= 0 or args.large_iterations <= 0: + raise ValueError("warmup must be non-negative and iteration counts must be positive") + + dtype = getattr(torch, args.dtype) + torch.manual_seed(0) + print_header(args, dtype) + with torch.inference_mode(): + for m in (int(value) for value in args.m.split(",")): + if m <= 0: + raise ValueError(f"M must be positive, got {m}") + benchmark_shape(m, args, dtype) + + +if __name__ == "__main__": + main() diff --git a/tests/diffusion/lora/test_base_linear.py b/tests/diffusion/lora/test_base_linear.py index b2f76be05..f3ecf6901 100644 --- a/tests/diffusion/lora/test_base_linear.py +++ b/tests/diffusion/lora/test_base_linear.py @@ -29,7 +29,7 @@ class _DummyQuantMethod: return y -def test_diffusion_base_linear_apply_multi_slice(): +def test_diffusion_base_linear_apply_multi_slice(monkeypatch): # Build a fake diffusion LoRA layer with 2 slices and rank=2. layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA) layer.tp_size = 1 @@ -72,8 +72,18 @@ def test_diffusion_base_linear_apply_multi_slice(): layer.lora_b_stacked = (b0, b1) layer.output_slices = out_slices + addmm_calls = [] + original_addmm = torch.addmm + + def _record_addmm(input, mat1, mat2, *, out=None): + addmm_calls.append((input, out)) + return original_addmm(input, mat1, mat2, out=out) + + monkeypatch.setattr(torch, "addmm", _record_addmm) + x = torch.tensor([[1.0, 2.0, 3.0]]) - out = layer.apply(x) + with torch.inference_mode(): + out = layer.apply(x) # Base output is identity: [1,2,3] base_out = x @ base_weight.t() @@ -87,6 +97,66 @@ def test_diffusion_base_linear_apply_multi_slice(): delta1 = torch.tensor([[6.0]]) expected = torch.cat([base_out[:, :2] + delta0, base_out[:, 2:3] + delta1], dim=-1) assert torch.allclose(out, expected) + assert len(addmm_calls) == 2 + assert all(input is out for input, out in addmm_calls) + + +def test_diffusion_base_linear_apply_preserves_autograd(monkeypatch): + layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA) + layer.tp_size = 1 + layer.lora_config = _DummyLoRAConfig() + + layer.base_layer = type("Base", (), {})() + layer.base_layer.quant_method = _DummyQuantMethod(torch.eye(2)) + + a = torch.tensor([[[[1.0, 2.0]]]], requires_grad=True) + b = torch.tensor([[[[3.0], [4.0]]]], requires_grad=True) + layer.lora_a_stacked = (a,) + layer.lora_b_stacked = (b,) + layer.output_slices = (2,) + layer._diffusion_lora_active_slices = (True,) + + def _unexpected_addmm(*args, **kwargs): + raise AssertionError("the out= fast path must not run with autograd enabled") + + monkeypatch.setattr(torch, "addmm", _unexpected_addmm) + + x = torch.tensor([[2.0, 1.0]], requires_grad=True) + out = layer.apply(x) + assert torch.allclose(out, torch.tensor([[14.0, 17.0]])) + + out.sum().backward() + assert torch.allclose(x.grad, torch.tensor([[8.0, 15.0]])) + assert torch.allclose(a.grad, torch.tensor([[[[14.0, 7.0]]]])) + assert torch.allclose(b.grad, torch.tensor([[[[4.0], [4.0]]]])) + + +def test_diffusion_base_linear_apply_falls_back_for_mixed_output_dtype(monkeypatch): + layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA) + layer.tp_size = 1 + layer.lora_config = _DummyLoRAConfig() + + class _Float64OutputQuantMethod(_DummyQuantMethod): + def apply(self, base_layer, x: torch.Tensor, bias: torch.Tensor | None): + return super().apply(base_layer, x, bias).to(torch.float64) + + layer.base_layer = type("Base", (), {})() + layer.base_layer.quant_method = _Float64OutputQuantMethod(torch.eye(2)) + layer.lora_a_stacked = (torch.ones((1, 1, 1, 2)),) + layer.lora_b_stacked = (torch.tensor([[[[1.0], [2.0]]]]),) + layer.output_slices = (2,) + layer._diffusion_lora_active_slices = (True,) + + def _unexpected_addmm(*args, **kwargs): + raise AssertionError("mixed output and LoRA dtypes must use the functional fallback") + + monkeypatch.setattr(torch, "addmm", _unexpected_addmm) + + with torch.inference_mode(): + out = layer.apply(torch.tensor([[1.0, 2.0]])) + + assert out.dtype == torch.float64 + assert torch.allclose(out, torch.tensor([[4.0, 8.0]], dtype=torch.float64)) def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch): diff --git a/vllm_omni/diffusion/lora/layers/base_linear.py b/vllm_omni/diffusion/lora/layers/base_linear.py index fae385c36..c1f941a10 100644 --- a/vllm_omni/diffusion/lora/layers/base_linear.py +++ b/vllm_omni/diffusion/lora/layers/base_linear.py @@ -143,8 +143,17 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA): # LoRA shrink & expand as in add_lora_linear(): # buffer = (x @ A.T) # y += buffer @ B.T - delta = (x_flat @ A.t()) @ B.t() - y_flat[:, offset : offset + slice_size] = y_flat[:, offset : offset + slice_size] + delta + buffer = x_flat @ A.t() + y_slice = y_flat[:, offset : offset + slice_size] + + # In inference, accumulate the expand GEMM directly into the base + # output. This avoids materializing both a full-width delta and the + # result of the following elementwise addition. ``out=`` does not + # support autograd, so training keeps the functional fallback. + if not torch.is_grad_enabled() and y_slice.dtype == buffer.dtype == B.dtype: + torch.addmm(y_slice, buffer, B.t(), out=y_slice) + else: + y_slice[:] = y_slice + buffer @ B.t() offset += slice_size return y_flat.view(original_shape)