BeatSeatandamy-why-3459
b308e196c9
Signed-off-by: BeatSeat <wendavid552@gmail.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
886 lines
41 KiB
Python
886 lines
41 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterable
|
|
from dataclasses import dataclass
|
|
from typing import ClassVar
|
|
|
|
import torch
|
|
from diffusers.image_processor import VaeImageProcessor
|
|
from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
from torch import nn
|
|
from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
|
|
|
|
from vllm_omni.diffusion.cache.cachedit import (
|
|
CacheDiTBackend,
|
|
CacheDiTRequestSpec,
|
|
RequestScopedCacheDiTRuntime,
|
|
)
|
|
from vllm_omni.diffusion.data import DiffusionCacheConfig, DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.utils import get_local_device
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config
|
|
|
|
from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DModel
|
|
from .rope_real import RotaryPosEmbedReal
|
|
from .schedulers import FlowMatchEulerDiscreteScheduler
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# Identifies the pipeline-owned Cache-DiT installation across requests.
|
|
_MAMMOTHMODA2_CACHE_DIT_KEY = "mammothmoda2:cache_dit"
|
|
|
|
|
|
def _first_request_value(value: object) -> object:
|
|
if isinstance(value, (list, tuple)):
|
|
return value[0] if value else None
|
|
return value
|
|
|
|
|
|
def get_mammoth_moda2_post_process_func(
|
|
_od_config: OmniDiffusionConfig,
|
|
):
|
|
image_processor = VaeImageProcessor()
|
|
|
|
def post_process_func(images: torch.Tensor):
|
|
return image_processor.postprocess(images)
|
|
|
|
return post_process_func
|
|
|
|
|
|
def _build_mammoth_config(od_config: OmniDiffusionConfig) -> Mammothmoda2Config:
|
|
raw_config = od_config.tf_model_config.to_dict()
|
|
if not raw_config:
|
|
raise ValueError("MammothModa2 diffusion stage requires the root checkpoint config")
|
|
return Mammothmoda2Config(**raw_config)
|
|
|
|
|
|
def _root_weight_source(
|
|
od_config: OmniDiffusionConfig,
|
|
) -> DiffusersPipelineLoader.ComponentSource:
|
|
if not od_config.model:
|
|
raise ValueError("MammothModa2 diffusion stage requires a model path")
|
|
return DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=od_config.model,
|
|
subfolder=None,
|
|
revision=od_config.revision,
|
|
prefix="",
|
|
fall_back_to_pt=True,
|
|
)
|
|
|
|
|
|
def _pad_cond_sequence(
|
|
embeds: list[torch.Tensor],
|
|
masks: list[torch.Tensor],
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Right-pad per-request ``[1, T_i, H]`` embeds into ``[B, T_max, H]``.
|
|
|
|
The transformer derives per-row ``encoder_seq_lengths`` from the boolean
|
|
mask, so padding must sit on the right and masked positions are excluded
|
|
from the packed joint sequence.
|
|
"""
|
|
if not embeds:
|
|
raise ValueError("Cannot pad an empty conditioning sequence list")
|
|
if len(embeds) != len(masks):
|
|
raise ValueError(f"Conditioning embeds/mask count mismatch: {len(embeds)} vs {len(masks)}")
|
|
|
|
# Fast-path for single-item batch: bypass padding tensor allocation
|
|
if len(embeds) == 1:
|
|
emb0, m0 = embeds[0], masks[0]
|
|
if emb0.ndim != 3 or emb0.shape[0] != 1:
|
|
raise ValueError(f"Conditioning embeds[0] must be [1, T, H], got {tuple(emb0.shape)}")
|
|
if m0.ndim != 2 or m0.shape != emb0.shape[:2]:
|
|
raise ValueError(f"Conditioning embeds/mask length mismatch: {tuple(emb0.shape)} vs {tuple(m0.shape)}")
|
|
return emb0, m0
|
|
|
|
hidden = embeds[0].shape[-1]
|
|
lengths = [int(e.shape[1]) for e in embeds]
|
|
max_len = max(lengths)
|
|
|
|
# Fast-path for homogeneous lengths: concatenate directly in C++
|
|
if all(length == max_len for length in lengths):
|
|
for i, (emb, m) in enumerate(zip(embeds, masks)):
|
|
if emb.ndim != 3 or emb.shape[0] != 1 or emb.shape[-1] != hidden:
|
|
raise ValueError(f"Conditioning embeds[{i}] must be [1, T, {hidden}], got {tuple(emb.shape)}")
|
|
if m.ndim != 2 or m.shape[0] != 1 or m.shape[1] != max_len:
|
|
raise ValueError(f"Conditioning embeds/mask length mismatch: {tuple(emb.shape)} vs {tuple(m.shape)}")
|
|
return torch.cat(embeds, dim=0), torch.cat(masks, dim=0)
|
|
|
|
out = embeds[0].new_zeros((len(embeds), max_len, hidden))
|
|
mask = torch.zeros((len(embeds), max_len), dtype=torch.bool, device=out.device)
|
|
for i, (emb, m) in enumerate(zip(embeds, masks)):
|
|
if emb.ndim != 3 or emb.shape[0] != 1 or emb.shape[-1] != hidden:
|
|
raise ValueError(f"Conditioning embeds[{i}] must be [1, T, {hidden}], got {tuple(emb.shape)}")
|
|
if m.ndim != 2 or m.shape[0] != 1 or m.shape[1] != emb.shape[1]:
|
|
raise ValueError(f"Conditioning embeds/mask length mismatch: {tuple(emb.shape)} vs {tuple(m.shape)}")
|
|
out[i, : emb.shape[1]] = emb[0]
|
|
mask[i, : m.shape[1]] = m[0].bool()
|
|
return out, mask
|
|
|
|
|
|
def _resolve_request_sampling_and_dims(
|
|
prompt: dict,
|
|
sampling: OmniDiffusionSamplingParams | None,
|
|
request_id: str,
|
|
) -> tuple[int, int, float, tuple[float, float], int]:
|
|
"""Shared extraction and validation of geometry and sampling knobs."""
|
|
dimensions = []
|
|
for name in ("height", "width"):
|
|
value = DiffusionRequestBatch.get_prompt_field(prompt, name)
|
|
if value is None and sampling is not None:
|
|
value = getattr(sampling, name, None)
|
|
if value is None:
|
|
value = 1024
|
|
try:
|
|
dimensions.append(int(value))
|
|
except (TypeError, ValueError, OverflowError) as exc:
|
|
raise ValueError(f"Invalid image size {name}={value!r} for request {request_id}") from exc
|
|
height, width = dimensions
|
|
if height <= 0 or width <= 0:
|
|
raise ValueError(f"Invalid image size: {height}x{width} for request {request_id}")
|
|
if height % 16 != 0 or width % 16 != 0:
|
|
raise ValueError(f"Image size must be multiples of 16, got {height}x{width} for request {request_id}")
|
|
|
|
info = prompt.get("additional_information")
|
|
request_info = info if isinstance(info, dict) else {}
|
|
extra_args = sampling.extra_args or {} if sampling else {}
|
|
|
|
guidance = extra_args.get("text_guidance_scale")
|
|
if guidance is None and sampling is not None:
|
|
guidance = sampling.guidance_scale if sampling.guidance_scale_provided else None
|
|
if guidance is None:
|
|
guidance = _first_request_value(request_info.get("text_guidance_scale"))
|
|
if guidance is None:
|
|
guidance = 9.0
|
|
try:
|
|
text_guidance_scale = float(guidance)
|
|
except (TypeError, ValueError, OverflowError) as exc:
|
|
raise ValueError(f"Invalid text_guidance_scale for request {request_id}") from exc
|
|
|
|
raw_num_inference_steps = extra_args.get("num_inference_steps")
|
|
if raw_num_inference_steps is None and sampling is not None:
|
|
raw_num_inference_steps = sampling.num_inference_steps
|
|
if raw_num_inference_steps is None:
|
|
raw_num_inference_steps = _first_request_value(request_info.get("num_inference_steps"))
|
|
if raw_num_inference_steps is None:
|
|
raw_num_inference_steps = 50
|
|
try:
|
|
num_inference_steps = int(raw_num_inference_steps)
|
|
except (TypeError, ValueError, OverflowError) as exc:
|
|
raise ValueError(f"Invalid num_inference_steps for request {request_id}") from exc
|
|
if num_inference_steps <= 0:
|
|
raise ValueError(f"num_inference_steps must be positive for request {request_id}")
|
|
|
|
cfg_range = extra_args.get("cfg_range")
|
|
if cfg_range is None:
|
|
cfg_range = request_info.get("cfg_range")
|
|
if cfg_range is None:
|
|
cfg_range = [0.0, 1.0]
|
|
if not isinstance(cfg_range, (list, tuple)) or len(cfg_range) != 2:
|
|
raise ValueError(f"cfg_range requires two values for request {request_id}")
|
|
try:
|
|
cfg_start, cfg_end = float(cfg_range[0]), float(cfg_range[1])
|
|
except (TypeError, ValueError, OverflowError) as exc:
|
|
raise ValueError(f"cfg_range requires two values convertible to floats for request {request_id}") from exc
|
|
if not 0 <= cfg_start <= cfg_end <= 1:
|
|
raise ValueError(f"cfg_range must satisfy 0 <= start <= end <= 1 for request {request_id}")
|
|
|
|
return height, width, text_guidance_scale, (cfg_start, cfg_end), num_inference_steps
|
|
|
|
|
|
def _resolve_gen_vocab_start_index(od_config: OmniDiffusionConfig | None) -> int:
|
|
"""Resolve the visual token threshold through the normalized model configuration."""
|
|
if od_config is not None and getattr(od_config, "tf_model_config", None) is not None:
|
|
try:
|
|
cfg = _build_mammoth_config(od_config)
|
|
llm_cfg = getattr(cfg, "llm_config", None)
|
|
if llm_cfg is not None:
|
|
val = getattr(llm_cfg, "gen_vocab_start_index", None)
|
|
if val is not None:
|
|
return int(val)
|
|
text_cfg = getattr(llm_cfg, "text_config", None)
|
|
if text_cfg is not None and getattr(text_cfg, "gen_vocab_start_index", None) is not None:
|
|
return int(text_cfg.gen_vocab_start_index)
|
|
except Exception:
|
|
pass
|
|
|
|
# Fallback inspection of raw dictionary if _build_mammoth_config fails (e.g. mock test objects)
|
|
raw_cfg = od_config.tf_model_config.to_dict()
|
|
if isinstance(raw_cfg, dict):
|
|
llm_cfg = raw_cfg.get("llm_config")
|
|
if isinstance(llm_cfg, dict):
|
|
text_cfg = llm_cfg.get("text_config")
|
|
if isinstance(text_cfg, dict) and text_cfg.get("gen_vocab_start_index") is not None:
|
|
return int(text_cfg["gen_vocab_start_index"])
|
|
if llm_cfg.get("gen_vocab_start_index") is not None:
|
|
return int(llm_cfg["gen_vocab_start_index"])
|
|
return 152064
|
|
|
|
|
|
def _validate_request_for_admission(
|
|
request: OmniDiffusionRequest,
|
|
od_config: OmniDiffusionConfig | None = None,
|
|
gen_vocab_start_index: int | None = None,
|
|
) -> tuple[int, int, int]:
|
|
"""Validate request at admission time to fail-fast before entering scheduler queue.
|
|
|
|
Returns (height, width, num_inference_steps).
|
|
"""
|
|
request_id = request.request_id
|
|
prompt = request.prompt if isinstance(request.prompt, dict) else {}
|
|
sampling = request.sampling_params
|
|
if sampling is not None and getattr(sampling, "num_outputs_per_prompt", 1) != 1:
|
|
raise ValueError(
|
|
f"MammothModa2 requires num_outputs_per_prompt == 1, got {sampling.num_outputs_per_prompt} "
|
|
f"for request {request_id}"
|
|
)
|
|
if not request.is_dummy_run():
|
|
info = prompt.get("additional_information")
|
|
if not isinstance(info, dict):
|
|
raise ValueError(f"Missing additional_information AR conditions for request {request_id}")
|
|
full_hidden_states = info.get("full_hidden_states")
|
|
full_token_ids = info.get("full_token_ids")
|
|
if not isinstance(full_hidden_states, torch.Tensor) or not isinstance(full_token_ids, list):
|
|
raise ValueError(f"Expected full_hidden_states tensor and full_token_ids list for request {request_id}")
|
|
try:
|
|
answer_start_index = int(info.get("answer_start_index"))
|
|
except (TypeError, ValueError, OverflowError) as exc:
|
|
raise ValueError(f"Invalid answer_start_index for request {request_id}") from exc
|
|
if full_hidden_states.ndim != 2:
|
|
raise ValueError(f"Expected 2D full_hidden_states for request {request_id}")
|
|
if full_hidden_states.shape[0] != len(full_token_ids):
|
|
raise ValueError(f"AR hidden-state/token-count mismatch for request {request_id}")
|
|
if not 0 <= answer_start_index <= len(full_token_ids):
|
|
raise ValueError(f"answer_start_index outside token range for request {request_id}")
|
|
try:
|
|
int_tokens = [int(token_id) for token_id in full_token_ids]
|
|
except (TypeError, ValueError, OverflowError) as exc:
|
|
raise ValueError(f"Invalid full_token_ids for request {request_id}") from exc
|
|
|
|
if gen_vocab_start_index is None:
|
|
gen_vocab_start_index = _resolve_gen_vocab_start_index(od_config)
|
|
answer_tokens = int_tokens[answer_start_index:]
|
|
if not any(token_id >= gen_vocab_start_index for token_id in answer_tokens):
|
|
raise ValueError(
|
|
f"MammothModa2 AR stage produced no visual-token hidden states for request {request_id}; "
|
|
f"the DiT stage requires at least one generated visual token. Generated token ids: {answer_tokens[:32]}"
|
|
)
|
|
|
|
height, width, _, _, num_inference_steps = _resolve_request_sampling_and_dims(prompt, sampling, request_id)
|
|
return height, width, num_inference_steps
|
|
|
|
|
|
def get_mammoth_moda2_pre_process_func(od_config: OmniDiffusionConfig | None = None):
|
|
"""Admission preprocessor: fail-fast per-request validation and grouping.
|
|
|
|
Validates AR conditions, dimensions, and sampling knobs at admission so
|
|
malformed requests are rejected individually with their request_id before
|
|
entering the scheduler queue. Admitted requests are assigned a
|
|
batch_compatibility_key based on output geometry and inference steps.
|
|
"""
|
|
gen_vocab_start_index = _resolve_gen_vocab_start_index(od_config)
|
|
|
|
def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
|
|
height, width, num_inference_steps = _validate_request_for_admission(
|
|
request, od_config=od_config, gen_vocab_start_index=gen_vocab_start_index
|
|
)
|
|
request.batch_compatibility_key = (
|
|
"mammoth_moda2_dit",
|
|
height,
|
|
width,
|
|
num_inference_steps,
|
|
)
|
|
return request
|
|
|
|
return pre_process_func
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _MammothRequest:
|
|
index: int
|
|
request_id: str
|
|
full_hidden_states: torch.Tensor
|
|
full_token_ids: list[int]
|
|
answer_start_index: int
|
|
height: int
|
|
width: int
|
|
text_guidance_scale: float
|
|
cfg_range: tuple[float, float]
|
|
num_inference_steps: int
|
|
seed: int | None
|
|
generator: torch.Generator | list[torch.Generator] | None
|
|
generator_device: torch.device | str | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _RequestConditioning:
|
|
text_embeds: torch.Tensor
|
|
text_mask: torch.Tensor
|
|
image_embeds: torch.Tensor
|
|
image_mask: torch.Tensor
|
|
|
|
|
|
class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
|
|
"""
|
|
MammothModa2 DiT + VAE generation stage (non-autoregressive).
|
|
|
|
This stage expects "image condition token hidden states" from the upstream AR stage,
|
|
and outputs image tensors via diffusion transformer + VAE decode.
|
|
|
|
"""
|
|
|
|
_dit_modules: ClassVar[list[str]] = ["gen_transformer"]
|
|
_encoder_modules: ClassVar[list[str]] = ["gen_image_condition_refiner"]
|
|
_vae_modules: ClassVar[list[str]] = ["gen_vae"]
|
|
|
|
supports_request_batch = True
|
|
supports_step_execution = False
|
|
|
|
# Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage.
|
|
hf_to_vllm_mapper = WeightsMapper(
|
|
orig_to_new_prefix={
|
|
"llm_model.": None,
|
|
"gen_tokenizer.": None,
|
|
}
|
|
)
|
|
|
|
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
|
|
super().__init__()
|
|
del prefix
|
|
self.od_config = od_config
|
|
self.device = get_local_device()
|
|
self.config = _build_mammoth_config(od_config)
|
|
self.weights_sources = [_root_weight_source(od_config)]
|
|
|
|
# --- Build DiT / VAE modules (names must match checkpoint keys) ---
|
|
if self.config.gen_vae_config is None or self.config.gen_dit_config is None:
|
|
raise ValueError("Mammothmoda2Config.gen_vae_config / gen_dit_config must not be None")
|
|
|
|
self.gen_vae = AutoencoderKL.from_config(self.config.gen_vae_config)
|
|
self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config)
|
|
|
|
# llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config
|
|
llm_hidden_size = 0
|
|
text_config = self.config.get_text_config()
|
|
if text_config is None:
|
|
logger.warning("No text config; failed to infer llm_hidden_size.")
|
|
elif not hasattr(text_config, "hidden_size"):
|
|
logger.warning("Text config exists, but has no hidden_size attribute; failed to infer llm_hidden_size.")
|
|
else:
|
|
llm_hidden_size = int(text_config.hidden_size or 0)
|
|
if llm_hidden_size <= 0:
|
|
raise ValueError(
|
|
"Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.text_config.hidden_size"
|
|
)
|
|
self._reinit_caption_embedder(llm_hidden_size)
|
|
|
|
# Optional image condition Q-Former. Preview stores it as a standalone
|
|
# module; Dev stores it under the DiT timestep/caption embedder.
|
|
llm_model_type = getattr(self.config.llm_config, "model_type", "")
|
|
refiner_config = self.config.gen_image_condition_refiner_config
|
|
if refiner_config is not None and llm_model_type == "mammothmoda2_qwen3_vl":
|
|
dit_hidden_size = int(self.gen_transformer.hidden_size)
|
|
self.gen_transformer.time_caption_embed.image_embedder = SimpleQFormerImageRefiner(
|
|
hidden_size=llm_hidden_size,
|
|
output_hidden_size=dit_hidden_size,
|
|
num_heads=max(1, dit_hidden_size // 128),
|
|
**refiner_config,
|
|
)
|
|
self.gen_image_condition_refiner = None
|
|
elif refiner_config is not None:
|
|
self.gen_image_condition_refiner = SimpleQFormerImageRefiner(
|
|
hidden_size=llm_hidden_size,
|
|
**refiner_config,
|
|
)
|
|
else:
|
|
self.gen_image_condition_refiner = None
|
|
|
|
# Precompute rotary freqs for diffusion transformer
|
|
# IMPORTANT: follow upstream mammothmoda: use top-level `config.gen_axes_*`
|
|
# (the checkpoint's `gen_dit_config.axes_lens` can be as small as 1024,
|
|
# which is insufficient for vLLM dummy-run/cudagraph warmup).
|
|
self.gen_freqs_cis = RotaryPosEmbedReal.get_freqs_real(
|
|
tuple(self.config.gen_axes_dim_rope),
|
|
tuple(self.config.gen_axes_lens),
|
|
theta=10000,
|
|
)
|
|
|
|
self._llm_hidden_size = llm_hidden_size
|
|
|
|
# Cache-DiT lifecycle: the diffusion runner enables the configured
|
|
# backend (``cache_backend`` on the deploy YAML stage entry) at startup
|
|
# and transfers ownership here via the request-scoped protocol;
|
|
# forward() then reconciles per-request state (step count, CFG parity).
|
|
# The runner also emits the cache summary when
|
|
# ``enable_cache_dit_summary`` is set.
|
|
self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
|
|
self._cache_dit_config: DiffusionCacheConfig | None = None
|
|
if str(getattr(od_config, "cache_backend", "") or "").lower() == "cache_dit":
|
|
cache_config = od_config.cache_config
|
|
self._cache_dit_config = (
|
|
cache_config
|
|
if isinstance(cache_config, DiffusionCacheConfig)
|
|
else DiffusionCacheConfig.from_dict(cache_config or {})
|
|
)
|
|
|
|
def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
|
|
"""Adopt a runner-installed Cache-DiT backend (request-scoped protocol)."""
|
|
|
|
self._cache_dit_runtime.adopt(backend, installation_key=_MAMMOTHMODA2_CACHE_DIT_KEY)
|
|
|
|
def is_cache_dit_enabled(self) -> bool:
|
|
"""Return the request-scoped Cache-DiT installation state."""
|
|
|
|
return self._cache_dit_runtime.is_enabled
|
|
|
|
def _prepare_cache_dit_for_request(self, *, num_inference_steps: int, cfg_active: bool) -> None:
|
|
"""Reconcile Cache-DiT hooks with this request before the denoise loop.
|
|
|
|
cache-dit's separate-CFG accounting assumes exactly two transformer
|
|
forwards per denoise step (conditional then unconditional), so CFG
|
|
requests refresh the context (per-request step counts vary) while
|
|
no-CFG requests run with hooks disabled: their single forward per
|
|
step cannot be represented by the parity accounting.
|
|
"""
|
|
if self._cache_dit_config is None:
|
|
return
|
|
if not cfg_active:
|
|
self._cache_dit_runtime.prepare(None)
|
|
return
|
|
self._cache_dit_runtime.prepare(
|
|
CacheDiTRequestSpec(
|
|
installation_key=_MAMMOTHMODA2_CACHE_DIT_KEY,
|
|
cache_config=self._cache_dit_config,
|
|
num_inference_steps=num_inference_steps,
|
|
)
|
|
)
|
|
|
|
def _reinit_caption_embedder(self, in_features: int) -> None:
|
|
# Align with upstream Mammothmoda2Model's `reinit_caption_embedder`:
|
|
# Use Qwen2RMSNorm(in_features) + Linear(in_features -> out_features).
|
|
out_features = int(getattr(self.gen_transformer, "hidden_size", 0) or self.gen_transformer.config.hidden_size)
|
|
self.gen_transformer.time_caption_embed.caption_embedder = nn.Sequential(
|
|
Qwen2RMSNorm(in_features, eps=1e-5),
|
|
nn.Linear(in_features, out_features, bias=True),
|
|
)
|
|
|
|
@staticmethod
|
|
def _group_requests(specs: list[_MammothRequest]) -> list[list[int]]:
|
|
"""Group co-scheduled requests by (height, width, num_inference_steps).
|
|
|
|
Preserves arrival order: the first-seen request of each bucket
|
|
determines the group's execution order.
|
|
"""
|
|
buckets: dict[tuple[int, int, int], list[int]] = {}
|
|
for idx, spec in enumerate(specs):
|
|
key = (spec.height, spec.width, spec.num_inference_steps)
|
|
buckets.setdefault(key, []).append(idx)
|
|
return list(buckets.values())
|
|
|
|
def _parse_request(self, req: DiffusionRequestBatch, index: int = 0) -> _MammothRequest:
|
|
request = req.requests[index]
|
|
request_id = request.request_id
|
|
prompt = request.prompt if isinstance(request.prompt, dict) else {}
|
|
sampling = request.sampling_params
|
|
if sampling is not None and getattr(sampling, "num_outputs_per_prompt", 1) != 1:
|
|
raise ValueError(
|
|
f"MammothModa2 requires num_outputs_per_prompt == 1, got {sampling.num_outputs_per_prompt} "
|
|
f"for request {request_id}"
|
|
)
|
|
info = prompt.get("additional_information")
|
|
if request.is_dummy_run():
|
|
gen_start = 152064
|
|
if hasattr(self, "config") and hasattr(self.config, "llm_config") and self.config.llm_config is not None:
|
|
val = getattr(self.config.llm_config, "gen_vocab_start_index", None)
|
|
if val is not None:
|
|
gen_start = val
|
|
else:
|
|
text_cfg = getattr(self.config.llm_config, "text_config", None)
|
|
if text_cfg is not None and getattr(text_cfg, "gen_vocab_start_index", None) is not None:
|
|
gen_start = text_cfg.gen_vocab_start_index
|
|
full_hidden_states = torch.zeros((2, self._llm_hidden_size), dtype=torch.float32)
|
|
full_token_ids = [0, int(gen_start)]
|
|
answer_start_index = 1
|
|
else:
|
|
if not isinstance(info, dict):
|
|
raise ValueError(f"Missing additional_information AR conditions for request {request_id}")
|
|
full_hidden_states = info.get("full_hidden_states")
|
|
full_token_ids = info.get("full_token_ids")
|
|
if not isinstance(full_hidden_states, torch.Tensor) or not isinstance(full_token_ids, list):
|
|
raise ValueError(f"Expected full_hidden_states tensor and full_token_ids list for request {request_id}")
|
|
try:
|
|
answer_start_index = int(info.get("answer_start_index"))
|
|
except (TypeError, ValueError, OverflowError) as exc:
|
|
raise ValueError(f"Invalid answer_start_index for request {request_id}") from exc
|
|
if full_hidden_states.ndim != 2:
|
|
raise ValueError(f"Expected 2D full_hidden_states for request {request_id}")
|
|
if full_hidden_states.shape[0] != len(full_token_ids):
|
|
raise ValueError(f"AR hidden-state/token-count mismatch for request {request_id}")
|
|
if not 0 <= answer_start_index <= len(full_token_ids):
|
|
raise ValueError(f"answer_start_index outside token range for request {request_id}")
|
|
try:
|
|
full_token_ids = [int(token_id) for token_id in full_token_ids]
|
|
except (TypeError, ValueError, OverflowError) as exc:
|
|
raise ValueError(f"Invalid full_token_ids for request {request_id}") from exc
|
|
|
|
height, width, text_guidance_scale, cfg_range, num_inference_steps = _resolve_request_sampling_and_dims(
|
|
prompt, sampling, request_id
|
|
)
|
|
|
|
generator = sampling.generator if sampling else None
|
|
if isinstance(generator, list) and len(generator) != 1:
|
|
raise ValueError(
|
|
f"MammothModa2 single-output request mode requires exactly one generator for request {request_id}"
|
|
)
|
|
|
|
return _MammothRequest(
|
|
index=index,
|
|
request_id=request_id,
|
|
full_hidden_states=full_hidden_states,
|
|
full_token_ids=full_token_ids,
|
|
answer_start_index=answer_start_index,
|
|
height=height,
|
|
width=width,
|
|
text_guidance_scale=text_guidance_scale,
|
|
cfg_range=cfg_range,
|
|
num_inference_steps=num_inference_steps,
|
|
seed=sampling.seed if sampling else None,
|
|
generator=generator,
|
|
generator_device=getattr(sampling, "generator_device", None) if sampling else None,
|
|
)
|
|
|
|
def _split_ar_conditions(
|
|
self,
|
|
*,
|
|
full_hidden_states: torch.Tensor,
|
|
full_token_ids: list[int],
|
|
answer_start_index: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Split AR-stage hidden states into text / image condition embeds.
|
|
|
|
The token ids that distinguish question (text) tokens, generated visual
|
|
tokens, and multi-modal placeholder tokens are read from the model config
|
|
(``gen_vocab_start_index`` and the vision placeholder token ids), so the
|
|
caller no longer needs to pass them. Mirrors the masking the bespoke
|
|
MammothModa2 example performed via ar2dit.
|
|
"""
|
|
gen_vocab_start_index = getattr(self.config.llm_config, "gen_vocab_start_index", None)
|
|
if gen_vocab_start_index is None:
|
|
text_cfg = getattr(self.config.llm_config, "text_config", None)
|
|
gen_vocab_start_index = getattr(text_cfg, "gen_vocab_start_index", 152064)
|
|
gen_vocab_start_index = int(gen_vocab_start_index)
|
|
cached_visual_ids = getattr(self, "_cached_visual_ids_tensor", None)
|
|
if cached_visual_ids is None:
|
|
cached_visual_ids = torch.tensor(
|
|
[
|
|
int(self.config.image_token_id),
|
|
int(self.config.video_token_id),
|
|
int(self.config.vision_start_token_id),
|
|
int(self.config.vision_end_token_id),
|
|
],
|
|
dtype=torch.long,
|
|
)
|
|
self._cached_visual_ids_tensor = cached_visual_ids
|
|
|
|
device = full_hidden_states.device
|
|
token_ids = torch.tensor(full_token_ids, dtype=torch.long, device=device)
|
|
positions = torch.arange(token_ids.shape[0], device=device)
|
|
questions_mask = positions < answer_start_index
|
|
answers_mask = ~questions_mask
|
|
gen_token_mask = token_ids >= gen_vocab_start_index
|
|
visual_ids_device = cached_visual_ids.to(device=device, non_blocking=True)
|
|
visual_token_mask = torch.isin(token_ids, visual_ids_device)
|
|
text_mask = questions_mask & ~(visual_token_mask | gen_token_mask)
|
|
image_mask = answers_mask & gen_token_mask
|
|
|
|
# Keep the transferred representation compact until the final device/dtype
|
|
# conversion in ``forward``. The masks are dtype-independent, and an
|
|
# unconditional host-side float32 expansion here would otherwise double both
|
|
# the selected-condition staging footprint and BF16 H2D traffic.
|
|
text_cond = full_hidden_states[text_mask].contiguous()
|
|
image_cond = full_hidden_states[image_mask].contiguous()
|
|
return text_cond, image_cond
|
|
|
|
def _split_request_conditions(self, request: _MammothRequest) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Split AR hidden states and fail fast before touching model params."""
|
|
text_cond, image_cond = self._split_ar_conditions(
|
|
full_hidden_states=request.full_hidden_states,
|
|
full_token_ids=request.full_token_ids,
|
|
answer_start_index=request.answer_start_index,
|
|
)
|
|
if image_cond.shape[0] == 0:
|
|
answer_token_ids = request.full_token_ids[request.answer_start_index :]
|
|
raise ValueError(
|
|
"MammothModa2 AR stage produced no visual-token hidden states for "
|
|
f"request {request.request_id}; the DiT stage requires at least one generated visual token. "
|
|
f"Generated token ids: {answer_token_ids[:32]}"
|
|
)
|
|
return text_cond, image_cond
|
|
|
|
def _build_conditioning(
|
|
self,
|
|
text_cond: torch.Tensor,
|
|
image_cond: torch.Tensor,
|
|
model_device: torch.device,
|
|
target_dtype: torch.dtype,
|
|
) -> _RequestConditioning:
|
|
text_embeds = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous().unsqueeze(0)
|
|
text_mask = torch.ones((1, text_embeds.shape[1]), dtype=torch.bool, device=model_device)
|
|
image_embeds = (
|
|
image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous().unsqueeze(0)
|
|
)
|
|
image_mask = torch.ones((1, image_embeds.shape[1]), dtype=torch.bool, device=model_device)
|
|
return _RequestConditioning(
|
|
text_embeds=text_embeds,
|
|
text_mask=text_mask,
|
|
image_embeds=image_embeds,
|
|
image_mask=image_mask,
|
|
)
|
|
|
|
@staticmethod
|
|
def _make_latent_generators(
|
|
specs: list[_MammothRequest],
|
|
fallback_device: torch.device,
|
|
) -> list[torch.Generator] | None:
|
|
generators: list[torch.Generator] = []
|
|
for s in specs:
|
|
if s.generator is not None:
|
|
gen = s.generator[0] if isinstance(s.generator, list) else s.generator
|
|
elif s.seed is not None:
|
|
device = s.generator_device or fallback_device
|
|
gen = torch.Generator(device=device).manual_seed(s.seed)
|
|
else:
|
|
return None
|
|
generators.append(gen)
|
|
return generators
|
|
|
|
def _denoise_group(
|
|
self,
|
|
specs: list[_MammothRequest],
|
|
conds: list[_RequestConditioning],
|
|
model_device: torch.device,
|
|
) -> list[torch.Tensor]:
|
|
batch = len(specs)
|
|
height = specs[0].height
|
|
width = specs[0].width
|
|
num_inference_steps = specs[0].num_inference_steps
|
|
target_dtype = conds[0].text_embeds.dtype
|
|
|
|
# Conditioning collation. Token order is text-then-image per request,
|
|
# so concatenate first and right-pad the combined sequence.
|
|
nested_image_embedder = getattr(self.gen_transformer.time_caption_embed, "image_embedder", None)
|
|
if self.gen_image_condition_refiner is not None:
|
|
image_embeds, image_mask = _pad_cond_sequence(
|
|
[c.image_embeds for c in conds],
|
|
[c.image_mask for c in conds],
|
|
)
|
|
# Apply optional refiner ONLY on image condition tokens.
|
|
if image_embeds.shape[1] > 0:
|
|
image_embeds = self.gen_image_condition_refiner(image_embeds, ~image_mask.bool())
|
|
refined_image_mask = torch.ones(
|
|
image_embeds.shape[:2],
|
|
dtype=torch.bool,
|
|
device=image_embeds.device,
|
|
)
|
|
seq_embeds = []
|
|
seq_masks = []
|
|
for i, c in enumerate(conds):
|
|
cur_img_embed = image_embeds[i : i + 1]
|
|
cur_img_mask = refined_image_mask[i : i + 1]
|
|
seq_embeds.append(torch.cat([c.text_embeds, cur_img_embed], dim=1))
|
|
seq_masks.append(torch.cat([c.text_mask, cur_img_mask], dim=1))
|
|
prompt_embeds, prompt_attention_mask = _pad_cond_sequence(seq_embeds, seq_masks)
|
|
ar_image_embeds = None
|
|
ar_image_attention_mask = None
|
|
elif nested_image_embedder is not None:
|
|
prompt_embeds, prompt_attention_mask = _pad_cond_sequence(
|
|
[c.text_embeds for c in conds],
|
|
[c.text_mask for c in conds],
|
|
)
|
|
ar_image_embeds, ar_image_attention_mask = _pad_cond_sequence(
|
|
[c.image_embeds for c in conds],
|
|
[c.image_mask for c in conds],
|
|
)
|
|
else:
|
|
prompt_embeds, prompt_attention_mask = _pad_cond_sequence(
|
|
[torch.cat([c.text_embeds, c.image_embeds], dim=1) for c in conds],
|
|
[torch.cat([c.text_mask, c.image_mask], dim=1) for c in conds],
|
|
)
|
|
ar_image_embeds = None
|
|
ar_image_attention_mask = None
|
|
|
|
# Empty unconditional prompt for classifier-free guidance, shared by
|
|
# all rows; rows that never take the CFG branch select their cond
|
|
# prediction directly via torch.where below.
|
|
needs_uncond = any(s.text_guidance_scale > 1.0 for s in specs)
|
|
negative_prompt_embeds = None
|
|
negative_prompt_attention_mask = None
|
|
if needs_uncond:
|
|
hidden_size = int(prompt_embeds.shape[-1])
|
|
negative_prompt_embeds = torch.zeros(
|
|
(batch, 0, hidden_size),
|
|
dtype=target_dtype,
|
|
device=model_device,
|
|
)
|
|
negative_prompt_attention_mask = torch.zeros(
|
|
(batch, 0),
|
|
dtype=torch.bool,
|
|
device=model_device,
|
|
)
|
|
|
|
vae_scale_factor = 16
|
|
latent_channels = int(self.gen_transformer.config.in_channels)
|
|
shape = (batch, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor)
|
|
single_shape = (1, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor)
|
|
generators = self._make_latent_generators(specs, model_device)
|
|
if generators is not None:
|
|
noise_list = []
|
|
for gen in generators:
|
|
gen_device = gen.device if hasattr(gen, "device") else model_device
|
|
n = randn_tensor(single_shape, generator=gen, device=gen_device, dtype=target_dtype)
|
|
noise_list.append(n.to(device=model_device))
|
|
latents = torch.cat(noise_list, dim=0)
|
|
else:
|
|
latents = randn_tensor(shape, device=model_device, dtype=target_dtype)
|
|
|
|
scheduler = FlowMatchEulerDiscreteScheduler()
|
|
scheduler.set_timesteps(
|
|
num_inference_steps=num_inference_steps,
|
|
device=model_device,
|
|
num_tokens=latents.shape[-2] * latents.shape[-1],
|
|
)
|
|
|
|
scale_vec = torch.tensor(
|
|
[s.text_guidance_scale for s in specs],
|
|
device=model_device,
|
|
dtype=target_dtype,
|
|
).view(batch, 1, 1, 1)
|
|
cfg_specs = [(s.text_guidance_scale > 1.0, float(s.cfg_range[0]), float(s.cfg_range[1])) for s in specs]
|
|
|
|
total_steps = max(1, len(scheduler.timesteps))
|
|
any_cfg_active = any(s.text_guidance_scale > 1.0 for s in specs)
|
|
self._prepare_cache_dit_for_request(
|
|
num_inference_steps=total_steps,
|
|
cfg_active=any_cfg_active,
|
|
)
|
|
requires_paired_cfg = self._cache_dit_runtime.is_enabled and getattr(
|
|
self, "_cache_dit_requires_paired_cfg", False
|
|
)
|
|
|
|
# Precompute mixed-CFG active masks across timesteps to avoid CPU list
|
|
# comprehensions and H2D copies in the step loop.
|
|
precomputed_active_masks = None
|
|
all_active_per_step = None
|
|
any_active_per_step = None
|
|
if needs_uncond:
|
|
precomputed_active_masks = torch.empty((total_steps, batch, 1, 1, 1), device=model_device, dtype=torch.bool)
|
|
all_active_per_step = []
|
|
any_active_per_step = []
|
|
for step_idx in range(total_steps):
|
|
frac = step_idx / total_steps
|
|
step_mask = [is_active and (lo <= frac <= hi) for is_active, lo, hi in cfg_specs]
|
|
all_active_per_step.append(all(step_mask))
|
|
any_active_per_step.append(any(step_mask))
|
|
mask_t = torch.tensor(step_mask, device=model_device, dtype=torch.bool)
|
|
precomputed_active_masks[step_idx] = mask_t.view(batch, 1, 1, 1)
|
|
|
|
# Precompute timesteps in latent dtype to eliminate per-step cast overhead
|
|
timesteps_dtype = scheduler.timesteps.to(latents.dtype)
|
|
|
|
# Run diffusion loop (CFG supported when text_guidance_scale > 1.0)
|
|
for i, t in enumerate(scheduler.timesteps):
|
|
timestep = timesteps_dtype[i].expand(batch)
|
|
model_pred = self.gen_transformer(
|
|
hidden_states=latents,
|
|
timestep=timestep,
|
|
text_hidden_states=prompt_embeds,
|
|
text_attention_mask=prompt_attention_mask,
|
|
ref_image_hidden_states=None,
|
|
ar_image_hidden_states=ar_image_embeds,
|
|
ar_image_attention_mask=ar_image_attention_mask,
|
|
freqs_cis=self.gen_freqs_cis,
|
|
)
|
|
run_uncond = (any_active_per_step is not None and any_active_per_step[i]) if needs_uncond else False
|
|
if requires_paired_cfg and negative_prompt_embeds is not None:
|
|
run_uncond = True
|
|
if run_uncond:
|
|
model_pred_uncond = self.gen_transformer(
|
|
hidden_states=latents,
|
|
timestep=timestep,
|
|
text_hidden_states=negative_prompt_embeds,
|
|
text_attention_mask=negative_prompt_attention_mask,
|
|
ref_image_hidden_states=None,
|
|
freqs_cis=self.gen_freqs_cis,
|
|
)
|
|
# Fused CFG blend: torch.lerp fuses (uncond + scale * (cond - uncond)) into 1 kernel
|
|
blended = torch.lerp(model_pred_uncond, model_pred, scale_vec)
|
|
if all_active_per_step is not None and all_active_per_step[i]:
|
|
# Fast path: all requests in the batch are active.
|
|
model_pred = blended
|
|
else:
|
|
# Inactive rows keep their conditional prediction exactly:
|
|
# gating the blend with torch.where avoids propagating
|
|
# uncond-branch NaN/rounding into CFG-free rows.
|
|
active_tensor = precomputed_active_masks[i]
|
|
model_pred = torch.where(active_tensor, blended, model_pred)
|
|
latents = scheduler.step(model_pred, t, latents, return_dict=False)[0]
|
|
latents = latents.to(dtype=target_dtype)
|
|
|
|
# VAE decode (in-place scaling to reduce peak VRAM before decode)
|
|
if self.gen_vae.config.scaling_factor is not None:
|
|
latents.div_(self.gen_vae.config.scaling_factor)
|
|
if self.gen_vae.config.shift_factor is not None:
|
|
latents.add_(self.gen_vae.config.shift_factor)
|
|
vae_dtype = next(self.gen_vae.parameters()).dtype
|
|
image = self.gen_vae.decode(latents.to(dtype=vae_dtype), return_dict=False)[0] # [B, C, H, W]
|
|
return [image[i : i + 1] for i in range(batch)]
|
|
|
|
@torch.inference_mode()
|
|
def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
|
|
specs = [self._parse_request(req, i) for i in range(req.num_reqs)]
|
|
|
|
# Validate AR conditions before touching model parameters so malformed
|
|
# requests fail fast instead of surfacing device/dtype errors.
|
|
raw_conds = [self._split_request_conditions(spec) for spec in specs]
|
|
|
|
model_device = next(self.parameters()).device
|
|
if self.gen_image_condition_refiner is not None:
|
|
target_dtype = next(self.gen_image_condition_refiner.parameters()).dtype
|
|
else:
|
|
target_dtype = next(self.gen_transformer.parameters()).dtype
|
|
|
|
conds = [
|
|
self._build_conditioning(text_cond, image_cond, model_device, target_dtype)
|
|
for text_cond, image_cond in raw_conds
|
|
]
|
|
|
|
groups = self._group_requests(specs)
|
|
if len(groups) == 1 and len(groups[0]) == len(specs):
|
|
images = self._denoise_group(specs, conds, model_device)
|
|
return [DiffusionOutput(output=img) for img in images]
|
|
|
|
outputs: list[DiffusionOutput | None] = [None] * len(specs)
|
|
for group in groups:
|
|
images = self._denoise_group(
|
|
[specs[i] for i in group],
|
|
[conds[i] for i in group],
|
|
model_device,
|
|
)
|
|
for index, image in zip(group, images):
|
|
outputs[index] = DiffusionOutput(output=image)
|
|
if any(output is None for output in outputs):
|
|
raise RuntimeError("DiT batching produced no image for at least one scheduled request")
|
|
return [output for output in outputs if output is not None]
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
loader = AutoWeightsLoader(self)
|
|
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
|