349 lines
12 KiB
Python
349 lines
12 KiB
Python
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
from torch.nn.parameter import Parameter
|
|||
|
|
from vllm.distributed import (
|
|||
|
|
tensor_model_parallel_all_gather,
|
|||
|
|
)
|
|||
|
|
from vllm.logger import init_logger
|
|||
|
|
from vllm.model_executor.layers.linear import (
|
|||
|
|
WEIGHT_LOADER_V2_SUPPORTED,
|
|||
|
|
QKVParallelLinear,
|
|||
|
|
)
|
|||
|
|
from vllm.model_executor.layers.quantization.base_config import (
|
|||
|
|
QuantizationConfig,
|
|||
|
|
)
|
|||
|
|
from vllm.model_executor.utils import set_weight_attrs
|
|||
|
|
from vllm.platforms import current_platform
|
|||
|
|
|
|||
|
|
from vllm_omni.diffusion.layers.mot.ops.mot_gemm import invoke_mot_gemm
|
|||
|
|
|
|||
|
|
logger = init_logger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MoTQKVParallelLinear(QKVParallelLinear):
|
|||
|
|
"""QKVParallelLinear with Mixture-of-Tokens routing.
|
|||
|
|
|
|||
|
|
Text weights: stored directly on self (self.weight, self.weight_scale, ...),
|
|||
|
|
created through the standard QKVParallelLinear.__init__ process.
|
|||
|
|
|
|||
|
|
VAE weights: stored in the permanent submodule self.gen_exp
|
|||
|
|
(self.gen_exp.weight, ...),
|
|||
|
|
created via quant_method.create_weights(self.gen_exp, ...).
|
|||
|
|
gen_exp.quant_method points to the same quant_method, so that
|
|||
|
|
the vLLM framework’s process_weights_after_loading
|
|||
|
|
can automatically detect and process it.
|
|||
|
|
|
|||
|
|
Forward behavior:
|
|||
|
|
- und mode (text_indices is None): fully reuse super().forward()
|
|||
|
|
- gen mode: call the MoT fused GEMM kernel
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
hidden_size: int,
|
|||
|
|
head_size: int,
|
|||
|
|
total_num_heads: int,
|
|||
|
|
total_num_kv_heads: int | None = None,
|
|||
|
|
bias: bool = True,
|
|||
|
|
vae_bias: bool = False,
|
|||
|
|
skip_bias_add: bool = False,
|
|||
|
|
params_dtype: torch.dtype | None = None,
|
|||
|
|
quant_config: QuantizationConfig | None = None,
|
|||
|
|
prefix: str = "",
|
|||
|
|
*,
|
|||
|
|
return_bias: bool = True,
|
|||
|
|
disable_tp: bool = False,
|
|||
|
|
v_head_size: int | None = None,
|
|||
|
|
):
|
|||
|
|
# ---- 1) Parent class creates text weights ----
|
|||
|
|
# super().__init__ will do the following:
|
|||
|
|
# quant_method.create_weights(self, ...) → self.weight, self.weight_scale, ...
|
|||
|
|
# QKVParallelLinear hardcodes gather_output=False
|
|||
|
|
super().__init__(
|
|||
|
|
hidden_size,
|
|||
|
|
head_size,
|
|||
|
|
total_num_heads,
|
|||
|
|
total_num_kv_heads,
|
|||
|
|
bias,
|
|||
|
|
skip_bias_add,
|
|||
|
|
params_dtype,
|
|||
|
|
quant_config,
|
|||
|
|
prefix,
|
|||
|
|
return_bias=return_bias,
|
|||
|
|
disable_tp=disable_tp,
|
|||
|
|
v_head_size=v_head_size,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# ---- 2) Create vae weights (permanent submodule) ----
|
|||
|
|
if self.quant_method is None:
|
|||
|
|
raise ValueError(
|
|||
|
|
f"quant_method must not be None for MoTQKVParallelLinear (prefix={prefix!r}). "
|
|||
|
|
"Ensure a valid QuantizationConfig is provided or the default UnquantizedLinearMethod is set."
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# NOTE: We instantiate a bare torch.nn.Module() here as a lightweight namespace
|
|||
|
|
# container to hold the secondary VAE weights.
|
|||
|
|
# This is a design choice for dynamically attaching `quant_method` to the bare module.
|
|||
|
|
# This is slightly unconventional but necessary. It allows vLLM's framework
|
|||
|
|
# (specifically `process_weights_after_loading`) to automatically discover
|
|||
|
|
# and process these secondary weights without requiring a dedicated,
|
|||
|
|
# boilerplate Module subclass just for the VAE pathway.
|
|||
|
|
self.gen_exp = torch.nn.Module()
|
|||
|
|
|
|||
|
|
# Use the same weight_loader as text
|
|||
|
|
vae_weight_loader = (
|
|||
|
|
self.weight_loader_v2
|
|||
|
|
if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED
|
|||
|
|
else self.weight_loader
|
|||
|
|
)
|
|||
|
|
self.quant_method.create_weights(
|
|||
|
|
layer=self.gen_exp,
|
|||
|
|
input_size_per_partition=self.input_size_per_partition,
|
|||
|
|
output_partition_sizes=self.output_partition_sizes,
|
|||
|
|
input_size=self.input_size,
|
|||
|
|
output_size=self.output_size,
|
|||
|
|
params_dtype=self.params_dtype,
|
|||
|
|
weight_loader=vae_weight_loader,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Make gen_exp discoverable by vLLM framework's process_weights_after_loading
|
|||
|
|
self.gen_exp.quant_method = self.quant_method
|
|||
|
|
|
|||
|
|
# ---- 3) vae bias ----
|
|||
|
|
if vae_bias:
|
|||
|
|
self.gen_exp.bias = Parameter(torch.empty(self.output_size_per_partition, dtype=self.params_dtype))
|
|||
|
|
set_weight_attrs(
|
|||
|
|
self.gen_exp.bias,
|
|||
|
|
{"output_dim": 0, "weight_loader": self.weight_loader},
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
self.gen_exp.register_parameter("bias", None)
|
|||
|
|
|
|||
|
|
self.update_param_tp_status()
|
|||
|
|
|
|||
|
|
# ==================================================================
|
|||
|
|
# Forward
|
|||
|
|
# ==================================================================
|
|||
|
|
def forward(
|
|||
|
|
self,
|
|||
|
|
input_: torch.Tensor,
|
|||
|
|
text_indices: torch.Tensor | None = None,
|
|||
|
|
vae_indices: torch.Tensor | None = None,
|
|||
|
|
) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
|
|||
|
|
# ---- und mode: fully reuse parent class (only text path) ----
|
|||
|
|
if text_indices is None:
|
|||
|
|
return super().forward(input_)
|
|||
|
|
|
|||
|
|
# ---- gen mode: fuse MoT GEMM ----
|
|||
|
|
output_parallel = self._mot_gemm_dispatch(input_, text_indices, vae_indices)
|
|||
|
|
|
|||
|
|
# QKVParallelLinear hardcodes gather_output=False, this branch never enters;
|
|||
|
|
# retained for future subclass changes
|
|||
|
|
if self.gather_output and self.tp_size > 1:
|
|||
|
|
output = tensor_model_parallel_all_gather(output_parallel)
|
|||
|
|
else:
|
|||
|
|
output = output_parallel
|
|||
|
|
|
|||
|
|
if not self.return_bias:
|
|||
|
|
return output
|
|||
|
|
|
|||
|
|
if self.skip_bias_add and text_indices is not None and (self.bias is not None or self.gen_exp.bias is not None):
|
|||
|
|
merged_bias = torch.zeros(
|
|||
|
|
output.size(0),
|
|||
|
|
self.output_size_per_partition,
|
|||
|
|
dtype=output.dtype,
|
|||
|
|
device=output.device,
|
|||
|
|
)
|
|||
|
|
if self.bias is not None:
|
|||
|
|
merged_bias[text_indices] = self.bias
|
|||
|
|
if self.gen_exp.bias is not None:
|
|||
|
|
merged_bias[vae_indices] = self.gen_exp.bias
|
|||
|
|
return output, merged_bias
|
|||
|
|
|
|||
|
|
output_bias = self.bias if self.skip_bias_add else None
|
|||
|
|
return output, output_bias
|
|||
|
|
|
|||
|
|
# ==================================================================
|
|||
|
|
# MoT GEMM dispatcher
|
|||
|
|
# ==================================================================
|
|||
|
|
def _mot_gemm_dispatch(
|
|||
|
|
self,
|
|||
|
|
x: torch.Tensor,
|
|||
|
|
text_indices: torch.Tensor,
|
|||
|
|
vae_indices: torch.Tensor,
|
|||
|
|
) -> torch.Tensor:
|
|||
|
|
# Any other backend (ROCm, XPU, TPU, CPU) uses the safe vllm fallback.
|
|||
|
|
if not current_platform.is_cuda():
|
|||
|
|
return self._mot_fallback(x, text_indices, vae_indices)
|
|||
|
|
|
|||
|
|
"""Dispatch to different MoT kernel paths based on weight dtype / quant attributes."""
|
|||
|
|
bias_text = self.bias if not self.skip_bias_add else None
|
|||
|
|
bias_vae = self.gen_exp.bias if not self.skip_bias_add else None
|
|||
|
|
|
|||
|
|
w_text = self.weight
|
|||
|
|
w_vae = self.gen_exp.weight
|
|||
|
|
assert w_text.dtype.is_floating_point == w_vae.dtype.is_floating_point, (
|
|||
|
|
"weight of text expert and image expert should be the same dtype."
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# w_text.dtype.itemsize >= 2 means bytes_per_element >= 2 (16bits or 32bits)
|
|||
|
|
if w_text.dtype.is_floating_point and w_text.dtype.itemsize >= 2:
|
|||
|
|
return self._mot_gemm_unquantized(
|
|||
|
|
x,
|
|||
|
|
text_indices,
|
|||
|
|
vae_indices,
|
|||
|
|
bias_text,
|
|||
|
|
bias_vae,
|
|||
|
|
)
|
|||
|
|
elif w_text.dtype == torch.float8_e4m3fn:
|
|||
|
|
return self._mot_gemm_fp8_w8a8(
|
|||
|
|
x,
|
|||
|
|
text_indices,
|
|||
|
|
vae_indices,
|
|||
|
|
bias_text,
|
|||
|
|
bias_vae,
|
|||
|
|
)
|
|||
|
|
elif w_text.dtype == torch.int8:
|
|||
|
|
return self._mot_gemm_weight_only(
|
|||
|
|
x,
|
|||
|
|
text_indices,
|
|||
|
|
vae_indices,
|
|||
|
|
bias_text,
|
|||
|
|
bias_vae,
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
return self._mot_fallback(
|
|||
|
|
x,
|
|||
|
|
text_indices,
|
|||
|
|
vae_indices,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# ==================================================================
|
|||
|
|
# Implementation of each quantization path
|
|||
|
|
# ==================================================================
|
|||
|
|
def _mot_gemm_unquantized(self, x, text_idx, vae_idx, bias_t, bias_v):
|
|||
|
|
"""BF16/FP16 path."""
|
|||
|
|
N = self.output_size_per_partition
|
|||
|
|
C = torch.zeros(x.size(0), N, dtype=x.dtype, device=x.device)
|
|||
|
|
invoke_mot_gemm(
|
|||
|
|
A=x,
|
|||
|
|
B_text=self.weight.data.t(),
|
|||
|
|
B_vae=self.gen_exp.weight.data.t(),
|
|||
|
|
C=C,
|
|||
|
|
bias_text=bias_t,
|
|||
|
|
bias_vae=bias_v,
|
|||
|
|
text_indices=text_idx,
|
|||
|
|
vae_indices=vae_idx,
|
|||
|
|
A_scale=None,
|
|||
|
|
B_text_scale=None,
|
|||
|
|
B_vae_scale=None,
|
|||
|
|
use_fp8_w8a8=False,
|
|||
|
|
use_int8_w8a8=False,
|
|||
|
|
use_int8_w8a16=False,
|
|||
|
|
use_int4_w4a16=False,
|
|||
|
|
A_per_channel_quant=False,
|
|||
|
|
B_per_channel_quant=False,
|
|||
|
|
)
|
|||
|
|
return C
|
|||
|
|
|
|||
|
|
def _mot_gemm_fp8_w8a8(self, x, text_idx, vae_idx, bias_t, bias_v):
|
|||
|
|
"""FP8 W8A8 path.
|
|||
|
|
|
|||
|
|
1) Activation quantization: reuse vllm's ops.scaled_fp8_quant
|
|||
|
|
2) MoT GEMM: text/vae use different fp8 weights and weight_scale
|
|||
|
|
3) De-quantization: done by MoT kernel's epilogue
|
|||
|
|
"""
|
|||
|
|
from vllm import _custom_ops as ops
|
|||
|
|
|
|||
|
|
x_2d = x.view(-1, x.shape[-1])
|
|||
|
|
input_scale = getattr(self, "input_scale", None)
|
|||
|
|
x_fp8, x_scale = ops.scaled_fp8_quant(
|
|||
|
|
x_2d,
|
|||
|
|
input_scale,
|
|||
|
|
use_per_token_if_dynamic=True,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
N = self.output_size_per_partition
|
|||
|
|
C = torch.zeros(x.size(0), N, dtype=x.dtype, device=x.device)
|
|||
|
|
|
|||
|
|
invoke_mot_gemm(
|
|||
|
|
A=x_fp8,
|
|||
|
|
B_text=self.weight.data,
|
|||
|
|
B_vae=self.gen_exp.weight.data,
|
|||
|
|
C=C,
|
|||
|
|
bias_text=bias_t,
|
|||
|
|
bias_vae=bias_v,
|
|||
|
|
text_indices=text_idx,
|
|||
|
|
vae_indices=vae_idx,
|
|||
|
|
A_scale=x_scale,
|
|||
|
|
B_text_scale=self.weight_scale.data,
|
|||
|
|
B_vae_scale=self.gen_exp.weight_scale.data,
|
|||
|
|
use_fp8_w8a8=True,
|
|||
|
|
use_int8_w8a8=False,
|
|||
|
|
use_int8_w8a16=False,
|
|||
|
|
use_int4_w4a16=False,
|
|||
|
|
A_per_channel_quant=True,
|
|||
|
|
B_per_channel_quant=False,
|
|||
|
|
)
|
|||
|
|
return C
|
|||
|
|
|
|||
|
|
def _mot_gemm_weight_only(self, x, text_idx, vae_idx, bias_t, bias_v):
|
|||
|
|
"""Weight-Only W8A16 path.
|
|||
|
|
|
|||
|
|
Activation values kept as bf16/fp16, weights are int8 + per-channel scale.
|
|||
|
|
De-quantization done by MoT kernel's epilogue.
|
|||
|
|
"""
|
|||
|
|
N = self.output_size_per_partition
|
|||
|
|
C = torch.zeros(x.size(0), N, dtype=x.dtype, device=x.device)
|
|||
|
|
invoke_mot_gemm(
|
|||
|
|
A=x,
|
|||
|
|
B_text=self.weight.data.t(),
|
|||
|
|
B_vae=self.gen_exp.weight.data.t(),
|
|||
|
|
C=C,
|
|||
|
|
bias_text=bias_t,
|
|||
|
|
bias_vae=bias_v,
|
|||
|
|
text_indices=text_idx,
|
|||
|
|
vae_indices=vae_idx,
|
|||
|
|
A_scale=None,
|
|||
|
|
B_text_scale=self.weight_scale.data,
|
|||
|
|
B_vae_scale=self.gen_exp.weight_scale.data,
|
|||
|
|
use_fp8_w8a8=False,
|
|||
|
|
use_int8_w8a8=False,
|
|||
|
|
use_int8_w8a16=True,
|
|||
|
|
use_int4_w4a16=False,
|
|||
|
|
A_per_channel_quant=False,
|
|||
|
|
B_per_channel_quant=True,
|
|||
|
|
)
|
|||
|
|
return C
|
|||
|
|
|
|||
|
|
def _mot_fallback(self, x, text_idx, vae_idx):
|
|||
|
|
"""Fallback: fall back to gather/scatter + quant_method.apply.
|
|||
|
|
|
|||
|
|
For unsupported quantization types, call standard forward for text/vae tokens.
|
|||
|
|
"""
|
|||
|
|
if self.quant_method is None:
|
|||
|
|
raise ValueError(
|
|||
|
|
"quant_method must not be None in MoTQKVParallelLinear._mot_fallback. "
|
|||
|
|
"This indicates an initialization error."
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
bias_text = self.bias if not self.skip_bias_add else None
|
|||
|
|
bias_vae = self.gen_exp.bias if not self.skip_bias_add else None
|
|||
|
|
|
|||
|
|
output = torch.zeros(
|
|||
|
|
x.size(0),
|
|||
|
|
self.output_size_per_partition,
|
|||
|
|
dtype=x.dtype,
|
|||
|
|
device=x.device,
|
|||
|
|
)
|
|||
|
|
output_text = self.quant_method.apply(self, x[text_idx], bias_text)
|
|||
|
|
|
|||
|
|
output_vae = self.quant_method.apply(
|
|||
|
|
self.gen_exp,
|
|||
|
|
x[vae_idx],
|
|||
|
|
bias_vae,
|
|||
|
|
)
|
|||
|
|
output[text_idx] = output_text
|
|||
|
|
output[vae_idx] = output_vae
|
|||
|
|
return output
|