355 lines
13 KiB
Python
355 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Quality-gated CUDA fast paths for Diffusers-style 2D KL VAE decoders."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import MethodType
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.vae_optimizations.gate import (
|
|
VaeFastPathGate,
|
|
register_vae_fast_path_gate,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
try:
|
|
from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import (
|
|
group_norm_silu_4d,
|
|
group_norm_silu_rows,
|
|
)
|
|
|
|
_HAS_TRITON = True
|
|
except ImportError: # pragma: no cover - exercised on builds without Triton
|
|
_HAS_TRITON = False
|
|
|
|
|
|
class FusedGroupNormSiLU(nn.Module):
|
|
"""Quality-gated GroupNorm + SiLU with an exact eager fallback."""
|
|
|
|
def __init__(self, norm: nn.GroupNorm, gate: VaeFastPathGate) -> None:
|
|
super().__init__()
|
|
self.num_groups = norm.num_groups
|
|
self.num_channels = norm.num_channels
|
|
self.eps = norm.eps
|
|
self.affine = norm.affine
|
|
# Register parameters directly so state-dict FQNs do not change.
|
|
self.weight = norm.weight
|
|
self.bias = norm.bias
|
|
self._vllm_gate = gate
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
if self._vllm_gate.enabled and x.dim() == 4 and not torch.compiler.is_compiling():
|
|
output = group_norm_silu_4d(
|
|
x,
|
|
self.weight,
|
|
self.bias,
|
|
self.num_groups,
|
|
self.eps,
|
|
apply_silu=True,
|
|
)
|
|
if output is not None:
|
|
return output
|
|
return F.silu(F.group_norm(x, self.num_groups, self.weight, self.bias, self.eps))
|
|
|
|
|
|
def _install_norm_silu(decoder: nn.Module, resnet_cls: type[nn.Module], gate: VaeFastPathGate) -> int:
|
|
def is_fusable(norm: object) -> bool:
|
|
return type(norm) is nn.GroupNorm and norm.affine and norm.weight is not None and norm.bias is not None
|
|
|
|
count = 0
|
|
for module in decoder.modules():
|
|
if (
|
|
type(module) is resnet_cls
|
|
and module.time_emb_proj is None
|
|
and module.time_embedding_norm in ("default", "group")
|
|
and module.upsample is None
|
|
and module.downsample is None
|
|
and isinstance(module.nonlinearity, nn.SiLU)
|
|
and is_fusable(module.norm1)
|
|
and is_fusable(module.norm2)
|
|
):
|
|
module.norm1 = FusedGroupNormSiLU(module.norm1, gate)
|
|
module.norm2 = FusedGroupNormSiLU(module.norm2, gate)
|
|
module.nonlinearity = nn.Identity()
|
|
count += 2
|
|
|
|
if (
|
|
type(getattr(decoder, "conv_norm_out", None)) is nn.GroupNorm
|
|
and isinstance(getattr(decoder, "conv_act", None), nn.SiLU)
|
|
and is_fusable(decoder.conv_norm_out)
|
|
):
|
|
decoder.conv_norm_out = FusedGroupNormSiLU(decoder.conv_norm_out, gate)
|
|
decoder.conv_act = nn.Identity()
|
|
count += 1
|
|
return count
|
|
|
|
|
|
_UPSAMPLE_TAP_MAP = {0: (2,), 1: (1, 2), 2: (0, 1), 3: (0,)}
|
|
|
|
|
|
def _fold_upsample2x_conv2d_weight(conv: nn.Conv2d) -> torch.Tensor:
|
|
weight = conv.weight.detach().float()
|
|
output_channels, input_channels = weight.shape[:2]
|
|
folded = weight.new_zeros(input_channels, output_channels, 4, 4)
|
|
for row in range(4):
|
|
for column in range(4):
|
|
value = weight.new_zeros(output_channels, input_channels)
|
|
for source_row in _UPSAMPLE_TAP_MAP[row]:
|
|
for source_column in _UPSAMPLE_TAP_MAP[column]:
|
|
value += weight[:, :, source_row, source_column]
|
|
folded[:, :, row, column] = value.t()
|
|
folded = folded.to(conv.weight.dtype)
|
|
if conv.weight.is_contiguous(memory_format=torch.channels_last):
|
|
folded = folded.contiguous(memory_format=torch.channels_last)
|
|
return folded.to(conv.weight.device)
|
|
|
|
|
|
class FusedUpsample2xConv2d(nn.Module):
|
|
"""Quality-gated nearest-2x + Conv2d replacement using ConvTranspose2d."""
|
|
|
|
def __init__(self, upsample: nn.Module, gate: VaeFastPathGate) -> None:
|
|
super().__init__()
|
|
# Keep only ``conv`` registered so checkpoint parameter names stay
|
|
# ``...upsamplers.N.conv.*``. The original module is an eager fallback.
|
|
object.__setattr__(self, "_original", upsample)
|
|
self.conv = upsample.conv
|
|
self.channels = upsample.channels
|
|
self._vllm_gate = gate
|
|
self.register_buffer("_fused_weight", None, persistent=False)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
output_size: int | None = None,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> torch.Tensor:
|
|
if not self._vllm_gate.enabled or output_size is not None or hidden_states.shape[1] != self.channels:
|
|
return self._original(hidden_states, output_size=output_size, *args, **kwargs)
|
|
|
|
weight = self._fused_weight
|
|
if weight is None or weight.device != self.conv.weight.device or weight.dtype != self.conv.weight.dtype:
|
|
weight = _fold_upsample2x_conv2d_weight(self.conv)
|
|
self._fused_weight = weight
|
|
return F.conv_transpose2d(hidden_states, weight, self.conv.bias, stride=2, padding=1)
|
|
|
|
|
|
def _install_fused_upsample(decoder: nn.Module, upsample_cls: type[nn.Module], gate: VaeFastPathGate) -> int:
|
|
count = 0
|
|
for block in decoder.up_blocks:
|
|
upsamplers = getattr(block, "upsamplers", None)
|
|
if not upsamplers:
|
|
continue
|
|
for index, upsample in enumerate(upsamplers):
|
|
if type(upsample) is not upsample_cls:
|
|
continue
|
|
conv = getattr(upsample, "conv", None)
|
|
if (
|
|
upsample.use_conv
|
|
and not upsample.use_conv_transpose
|
|
and upsample.interpolate
|
|
and upsample.norm is None
|
|
and upsample.name == "conv"
|
|
and type(conv) is nn.Conv2d
|
|
and conv.kernel_size == (3, 3)
|
|
and conv.stride == (1, 1)
|
|
and conv.padding == (1, 1)
|
|
and conv.dilation == (1, 1)
|
|
and conv.groups == 1
|
|
and conv.padding_mode == "zeros"
|
|
):
|
|
upsamplers[index] = FusedUpsample2xConv2d(upsample, gate)
|
|
count += 1
|
|
return count
|
|
|
|
|
|
def _fold_attention_value_projection(module: nn.Module) -> tuple[torch.Tensor, torch.Tensor]:
|
|
value_weight = module.to_v.weight.detach().float()
|
|
value_bias = module.to_v.bias.detach().float()
|
|
output_weight = module.to_out[0].weight.detach().float()
|
|
output_bias = module.to_out[0].bias.detach().float()
|
|
dtype = module.to_v.weight.dtype
|
|
return (
|
|
(output_weight @ value_weight).to(dtype).contiguous(),
|
|
(output_weight @ value_bias + output_bias).to(dtype),
|
|
)
|
|
|
|
|
|
def _attention_fast_forward(
|
|
self: nn.Module,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | None = None,
|
|
attention_mask: torch.Tensor | None = None,
|
|
temb: torch.Tensor | None = None,
|
|
**cross_attention_kwargs: Any,
|
|
) -> torch.Tensor:
|
|
if (
|
|
not self._vllm_gate.enabled
|
|
or encoder_hidden_states is not None
|
|
or attention_mask is not None
|
|
or temb is not None
|
|
or hidden_states.ndim != 4
|
|
):
|
|
return type(self).forward(
|
|
self,
|
|
hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
attention_mask=attention_mask,
|
|
temb=temb,
|
|
**cross_attention_kwargs,
|
|
)
|
|
|
|
residual = hidden_states
|
|
batch_size, channels, height, width = hidden_states.shape
|
|
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
|
|
|
|
if self.group_norm is not None:
|
|
norm = self.group_norm
|
|
normalized = None
|
|
if not torch.compiler.is_compiling():
|
|
normalized = group_norm_silu_rows(
|
|
hidden_states,
|
|
norm.weight,
|
|
norm.bias,
|
|
norm.num_groups,
|
|
norm.eps,
|
|
apply_silu=False,
|
|
)
|
|
hidden_states = normalized if normalized is not None else norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
|
|
|
query = self.to_q(hidden_states)
|
|
key = self.to_k(hidden_states)
|
|
folded_weight = self._vllm_folded_value_weight
|
|
folded_bias = self._vllm_folded_value_bias
|
|
if (
|
|
folded_weight is None
|
|
or folded_bias is None
|
|
or folded_weight.device != hidden_states.device
|
|
or folded_weight.dtype != self.to_v.weight.dtype
|
|
):
|
|
folded_weight, folded_bias = _fold_attention_value_projection(self)
|
|
self._vllm_folded_value_weight = folded_weight
|
|
self._vllm_folded_value_bias = folded_bias
|
|
value = F.linear(hidden_states, folded_weight, folded_bias)
|
|
|
|
output = F.scaled_dot_product_attention(query.unsqueeze(1), key.unsqueeze(1), value.unsqueeze(1))
|
|
output = output.squeeze(1).to(query.dtype)
|
|
output = self.to_out[1](output)
|
|
output = output.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
|
|
if self.residual_connection:
|
|
output = output + residual
|
|
return output / self.rescale_output_factor
|
|
|
|
|
|
def _attention_fast_compatible(module: nn.Module, attention_cls: type, processor_cls: type) -> bool:
|
|
return (
|
|
type(module) is attention_cls
|
|
and isinstance(module.processor, processor_cls)
|
|
and module.heads == 1
|
|
and module.scale_qk
|
|
and module.spatial_norm is None
|
|
and not module.norm_cross
|
|
and getattr(module, "norm_q", None) is None
|
|
and getattr(module, "norm_k", None) is None
|
|
and getattr(module, "add_k_proj", None) is None
|
|
and type(module.to_q) is nn.Linear
|
|
and type(module.to_k) is nn.Linear
|
|
and type(module.to_v) is nn.Linear
|
|
and type(module.to_out[0]) is nn.Linear
|
|
and isinstance(module.to_out[1], nn.Dropout)
|
|
and module.to_v.bias is not None
|
|
and module.to_out[0].bias is not None
|
|
)
|
|
|
|
|
|
def _decoder_layout_forward(self: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor:
|
|
use_channels_last = self._vllm_gate.enabled
|
|
if use_channels_last != self._vllm_channels_last:
|
|
memory_format = torch.channels_last if use_channels_last else torch.contiguous_format
|
|
self.to(memory_format=memory_format)
|
|
self._vllm_channels_last = use_channels_last
|
|
logger.debug(
|
|
"%s decoder switched to %s layout",
|
|
self._vllm_label,
|
|
"channels_last" if use_channels_last else "contiguous",
|
|
)
|
|
return type(self).forward(self, *args, **kwargs)
|
|
|
|
|
|
def maybe_optimize_image_vae(vae: nn.Module) -> nn.Module:
|
|
"""Install image VAE decoder fast paths when the structure is fully supported."""
|
|
|
|
if not _HAS_TRITON or getattr(vae, "_vllm_vae_fast_path_installed", False):
|
|
return vae
|
|
|
|
from diffusers.models.attention_processor import Attention, AttnProcessor2_0
|
|
from diffusers.models.autoencoders import AutoencoderKL
|
|
from diffusers.models.autoencoders.vae import Decoder
|
|
from diffusers.models.resnet import ResnetBlock2D
|
|
from diffusers.models.upsampling import Upsample2D
|
|
|
|
supported_types: tuple[type[nn.Module], ...] = (AutoencoderKL,)
|
|
try:
|
|
from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
|
|
|
|
supported_types += (AutoencoderKLFlux2,)
|
|
except ImportError: # pragma: no cover - older supported Diffusers builds
|
|
pass
|
|
|
|
if not isinstance(vae, supported_types):
|
|
return vae
|
|
decoder = getattr(vae, "decoder", None)
|
|
if type(decoder) is not Decoder:
|
|
return vae
|
|
|
|
attention_modules = [
|
|
module for module in decoder.modules() if _attention_fast_compatible(module, Attention, AttnProcessor2_0)
|
|
]
|
|
attention_count = sum(isinstance(module, Attention) for module in decoder.modules())
|
|
if len(attention_modules) != attention_count:
|
|
logger.info(
|
|
"%s has %d/%d attention blocks without a layout-safe rewrite; skipping VAE fast paths",
|
|
type(vae).__name__,
|
|
attention_count - len(attention_modules),
|
|
attention_count,
|
|
)
|
|
return vae
|
|
|
|
gate = VaeFastPathGate()
|
|
decoder._vllm_gate = gate
|
|
decoder._vllm_label = type(vae).__name__
|
|
decoder._vllm_channels_last = False
|
|
decoder.forward = MethodType(_decoder_layout_forward, decoder)
|
|
upsample_count = _install_fused_upsample(decoder, Upsample2D, gate)
|
|
for module in attention_modules:
|
|
module._vllm_gate = gate
|
|
module.register_buffer("_vllm_folded_value_weight", None, persistent=False)
|
|
module.register_buffer("_vllm_folded_value_bias", None, persistent=False)
|
|
module.forward = MethodType(_attention_fast_forward, module)
|
|
norm_count = _install_norm_silu(decoder, ResnetBlock2D, gate)
|
|
register_vae_fast_path_gate(vae, gate)
|
|
vae._vllm_vae_fast_path_installed = True
|
|
logger.info(
|
|
"%s installed quality-gated VAE decoder fast paths "
|
|
"(%d upsamplers, %d attention blocks, %d GroupNorm+SiLU fusions)",
|
|
type(vae).__name__,
|
|
upsample_count,
|
|
len(attention_modules),
|
|
norm_count,
|
|
)
|
|
return vae
|
|
|
|
|
|
__all__ = [
|
|
"FusedGroupNormSiLU",
|
|
"FusedUpsample2xConv2d",
|
|
"maybe_optimize_image_vae",
|
|
]
|