Skip to content

[diffusion] Add the dynamic LoRA delta inside the second GEMM #43501

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mickqian merged 3 commits into from
Oct 11, 2026
Merged

[diffusion] Add the dynamic LoRA delta inside the second GEMM #43501

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1198565
[diffusion] Add the dynamic LoRA delta inside the second GEMM
mickqian Oct 10, 2026
633fdf1
[diffusion] Scale the dynamic LoRA rows instead of the GEMM's alpha
mickqian Oct 10, 2026
ba4093b
[diffusion] Keep the exact LoRA chain for a forward without its request
mickqian Oct 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/docs/sglang-diffusion/fused_kernels.mdx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@ These fusion families mount under both `quality="lossless"` and
| LingBot Video fused RMSNorm | Replaces the handwritten cast, square, mean, rsqrt, and multiply chain with existing Triton RMSNorm kernels |
| LingBot Video per-token gated residual + RMSNorm modulate | Folds `residual + gate * update` (per-token `[B, S, 1]` gate) and the `rmsnorm(x) * (1 + scale) + shift` adaLN chain (strided `[B, S, 6D]` chunk views) into single kernels |
| SANA-Video BF16-input linear attention | Keeps the first linear-attention GEMM's inputs in BF16 with FP32 accumulation/output; the second GEMM remains FP32 |
| Dynamic LoRA delta | An unmerged adapter's scaled update accumulated onto the base output by its second GEMM (`addmm`, one rounding), and a row-parallel layer's bias added in the base GEMM's epilogue, instead of materializing the `[tokens, out]` delta and scaling and adding it in separate passes |
| FLUX-family VAE fast paths | Channels-last decode, GroupNorm(+SiLU), upsample, and attention replacements for FLUX.2 and AutoencoderKL-based FLUX.1, Z-Image, and SD3 pipelines |
| Wan VAE RMSNorm + SiLU | Replaces the channel-first RMSNorm/SiLU chain while keeping the decode in `channels_last_3d` |
| MiniMax-H3 VAE decoder fast path | Per-head fp32 RMSNorm, rounded to the activation dtype, and NeoX RoPE in one launch over the strided Q/K views of the ViT decoder's QKV buffer; cuDNN SDPA for the block's attention |
Expand Down
212 changes: 116 additions & 96 deletions python/sglang/multimodal_gen/runtime/layers/lora/linear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
)
from torch.distributed.tensor import DTensor

from sglang.multimodal_gen.configs.sample.sampling_params import quality_allows
from sglang.multimodal_gen.runtime.distributed import (
get_local_torch_device,
get_tp_rank,
Expand Down Expand Up @@ -76,6 +77,71 @@ def _compute_lora_delta(
)


def _request_allows_lossless() -> bool:
# only a request that allows it: many stages set the forward context
# without forward_batch, and an exact request must still get exact bytes
context = get_forward_context_or_none()
forward_batch = None if context is None else context.forward_batch
sampling_params = getattr(forward_batch, "sampling_params", None)
quality = getattr(sampling_params, "quality", None)
return quality is not None and quality_allows(quality, "lossless")


def _can_add_lora_delta_in_gemm(
out: torch.Tensor, lora_A: torch.Tensor, lora_B: torch.Tensor
) -> bool:
if not isinstance(lora_B, torch.Tensor) or lora_A.dim() != lora_B.dim():
return False
if lora_A.dim() == 3 and lora_A.shape[0] != lora_B.shape[0]:
return False
return lora_A.dim() in (2, 3) and out.is_contiguous() and _request_allows_lossless()


def _add_lora_delta_in_gemm(
out: torch.Tensor,
x: torch.Tensor,
lora_A: torch.Tensor,
lora_B: torch.Tensor,
scale: float,
) -> torch.Tensor:
"""``out + scale * delta``, accumulated onto ``out`` by the second GEMM.

Lossless tier: the reference path rounds the delta after the GEMM, after
each scale and after the add, moving the full ``[tokens, out]`` delta
through memory for each; here it is never materialized. The scale goes on
the rank-wide intermediate as a tensor op, not as the GEMM's ``alpha``: a
compiled forward reads the request scale as a symbolic float, which a
kwarg would freeze at its first value.
"""
out_2d = out.view(-1, out.shape[-1])
x_2d = x.reshape(-1, x.shape[-1])
if lora_A.dim() == 2:
hidden = x_2d @ lora_A.t()
else:
hidden = x_2d @ lora_A.reshape(-1, lora_A.shape[-1]).t()
if scale != 1.0:
hidden = hidden * scale
if lora_A.dim() == 2:
groups = [(out_2d, hidden, lora_B)]
else:
# stacked [groups, rank, in] / [groups, out, rank], one GEMM per group
rank, width = lora_A.shape[1], lora_B.shape[1]
groups = [
(
out_2d[:, g * width : (g + 1) * width],
hidden[:, g * rank : (g + 1) * rank],
lora_B[g],
)
for g in range(lora_B.shape[0])
]
for target, rows, b in groups:
if target.dtype == rows.dtype:
target.addmm_(rows, b.t())
else:
target.add_(rows @ b.t())
return out


class BaseLayerWithLoRA(nn.Module):
def __init__(
self,
Expand Down Expand Up @@ -154,34 +220,47 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:

@torch.compile()
def _forward_with_delta(self, x: torch.Tensor) -> torch.Tensor:
lora_A = self.lora_A
lora_B = self.lora_B
if isinstance(self.lora_B, DTensor):
lora_B = self.lora_B.to_local()
lora_A = self.lora_A.to_local()

# TODO: Support multiple LoRA adapters when use not merged mode
runtime_lora_scale = self._runtime_lora_scale()
if runtime_lora_scale == 0.0:
return self.base_layer(x)

lora_dtype = lora_A.dtype
x_lora = x.to(dtype=lora_dtype)
lora_A_sliced = self.slice_lora_a_weights(
out, output_bias = self.base_layer(x)
out = self._add_lora_delta(out, x, runtime_lora_scale)
return self._add_lora_output_offset(out), output_bias

def _add_lora_delta(
self, out: torch.Tensor, x: torch.Tensor, runtime_lora_scale: float
) -> torch.Tensor:
"""``out`` plus the active adapter's update for input ``x``."""
lora_A = self.lora_A
lora_B = self.lora_B
if isinstance(lora_B, DTensor):
lora_B = lora_B.to_local()
lora_A = lora_A.to_local()

x_lora = x.to(dtype=lora_A.dtype)
lora_A = self.slice_lora_a_weights(
lora_A.to(device=x.device, non_blocking=True)
)
lora_B_sliced = self.slice_lora_b_weights(
lora_B = self.slice_lora_b_weights(
lora_B.to(device=x.device, non_blocking=True)
)
delta = _compute_lora_delta(x_lora, lora_A_sliced, lora_B_sliced)
if _can_add_lora_delta_in_gemm(out, lora_A, lora_B):
scale = self.strength * runtime_lora_scale
if self.lora_alpha != self.lora_rank:
scale *= self.lora_alpha / self.lora_rank # type: ignore
return _add_lora_delta_in_gemm(out, x_lora, lora_A, lora_B, scale)

delta = _compute_lora_delta(x_lora, lora_A, lora_B)
if self.lora_alpha != self.lora_rank:
delta = delta * (
self.lora_alpha / self.lora_rank # type: ignore
) # type: ignore
delta = delta * self.strength * runtime_lora_scale
out, output_bias = self.base_layer(x)
out = out + delta.to(dtype=out.dtype)
return self._add_lora_output_offset(out), output_bias
delta = delta * (self.lora_alpha / self.lora_rank) # type: ignore
# x * 1.0 is exact, so skipping a unit scale keeps the reference bits
if self.strength != 1.0:
delta = delta * self.strength
if runtime_lora_scale != 1.0:
delta = delta * runtime_lora_scale
return out + delta.to(dtype=out.dtype)

def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
return A
Expand Down Expand Up @@ -657,35 +736,13 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor:
if runtime_lora_scale == 0.0:
return self.base_layer(input_)

lora_A = self.lora_A
lora_B = self.lora_B
if isinstance(self.lora_B, DTensor):
lora_B = self.lora_B.to_local()
lora_A = self.lora_A.to_local()

bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
output_parallel = self.base_layer.quant_method.apply(
self.base_layer, input_, bias
)
if not self.merged and not self.disable_lora:
lora_dtype = lora_A.dtype
input_lora = input_.to(dtype=lora_dtype)
lora_A_sliced = self.slice_lora_a_weights(
lora_A.to(device=input_.device, non_blocking=True)
)
lora_B_sliced = self.slice_lora_b_weights(
lora_B.to(device=input_.device, non_blocking=True)
)
delta_parallel = _compute_lora_delta(
input_lora, lora_A_sliced, lora_B_sliced
)
if self.lora_alpha != self.lora_rank:
delta_parallel = delta_parallel * (
self.lora_alpha / self.lora_rank # type: ignore
) # type: ignore
delta_parallel = delta_parallel * self.strength * runtime_lora_scale
output_parallel = output_parallel + delta_parallel.to(
dtype=output_parallel.dtype
output_parallel = self._add_lora_delta(
output_parallel, input_, runtime_lora_scale
)
output_parallel = self._add_lora_output_offset(output_parallel)
if self.base_layer.gather_output:
Expand Down Expand Up @@ -759,12 +816,6 @@ def forward(self, input_: torch.Tensor):
if runtime_lora_scale == 0.0:
return self.base_layer(input_)

lora_A = self.lora_A
lora_B = self.lora_B
if isinstance(self.lora_B, DTensor):
lora_B = self.lora_B.to_local()
lora_A = self.lora_A.to_local()

if self.base_layer.input_is_parallel:
input_parallel = input_
else:
Expand All @@ -773,28 +824,22 @@ def forward(self, input_: torch.Tensor):
input_, num_partitions=self.base_layer.tp_size
)
input_parallel = splitted_input[tp_rank].contiguous()
bias = self.base_layer.bias
# Lossless tier: the bias goes into the GEMM epilogue on rank 0, as in
# RowParallelLinear, instead of a separate pass after the reduce.
bias_in_gemm = (
bias is not None
and not self.base_layer.skip_bias_add
and _request_allows_lossless()
)
output_parallel = self.base_layer.quant_method.apply(
self.base_layer, input_parallel
self.base_layer,
input_parallel,
bias=bias if bias_in_gemm and self.base_layer.tp_rank == 0 else None,
)
if not self.merged and not self.disable_lora:
lora_dtype = lora_A.dtype
input_parallel_lora = input_parallel.to(dtype=lora_dtype)
lora_A_sliced = self.slice_lora_a_weights(
lora_A.to(device=input_parallel.device, non_blocking=True)
)
lora_B_sliced = self.slice_lora_b_weights(
lora_B.to(device=input_parallel.device, non_blocking=True)
)
delta_parallel = _compute_lora_delta(
input_parallel_lora, lora_A_sliced, lora_B_sliced
)
if self.lora_alpha != self.lora_rank:
delta_parallel = delta_parallel * (
self.lora_alpha / self.lora_rank # type: ignore
) # type: ignore
delta_parallel = delta_parallel * self.strength * runtime_lora_scale
output_parallel = output_parallel + delta_parallel.to(
dtype=output_parallel.dtype
output_parallel = self._add_lora_delta(
output_parallel, input_parallel, runtime_lora_scale
)

if self.base_layer.reduce_results and self.base_layer.tp_size > 1:
Expand All @@ -804,9 +849,7 @@ def forward(self, input_: torch.Tensor):

if not self.base_layer.skip_bias_add:
output = (
output_ + self.base_layer.bias
if self.base_layer.bias is not None
else output_
output_ + bias if bias is not None and not bias_in_gemm else output_
)
output_bias = None
else:
Expand All @@ -832,39 +875,16 @@ class LinearWithLoRA(BaseLayerWithLoRA):

@torch.compile()
def forward(self, x: torch.Tensor) -> torch.Tensor:
lora_A = self.lora_A
lora_B = self.lora_B
if isinstance(self.lora_B, DTensor):
lora_B = self.lora_B.to_local()
lora_A = self.lora_A.to_local()

# TODO: Support multiple LoRA adapters when use not merged mode
runtime_lora_scale = self._runtime_lora_scale()
if runtime_lora_scale == 0.0:
# nn.Linear.forward() returns a single tensor
return self.base_layer(x)

# nn.Linear.forward() returns a single tensor, not a tuple
out = self.base_layer(x)
if not self.merged and not self.disable_lora:
lora_dtype = lora_A.dtype
x_lora = x.to(dtype=lora_dtype)
lora_A_sliced = self.slice_lora_a_weights(
lora_A.to(device=x.device, non_blocking=True)
)
lora_B_sliced = self.slice_lora_b_weights(
lora_B.to(device=x.device, non_blocking=True)
)
delta = _compute_lora_delta(x_lora, lora_A_sliced, lora_B_sliced)
if self.lora_alpha != self.lora_rank:
delta = delta * (
self.lora_alpha / self.lora_rank # type: ignore
) # type: ignore
delta = delta * self.strength * runtime_lora_scale
# nn.Linear.forward() returns a single tensor, not a tuple
out = self.base_layer(x)
out = out + delta.to(dtype=out.dtype)
else:
# nn.Linear.forward() returns a single tensor
out = self.base_layer(x)
out = self._add_lora_delta(out, x, runtime_lora_scale)
return self._add_lora_output_offset(out)


Expand Down
Loading
Loading