wehub-resource-sync
1400 lines
57 KiB
Python
1400 lines
57 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""DiffusionGemma model, ModelState, and Sampler for vLLM.
|
|
|
|
Single Gemma4 backbone run in two modes (like YOCO):
|
|
- encoder mode: causal attention, writes KV cache
|
|
- decoder mode: bidirectional attention, reads encoder KV, doesn't write
|
|
|
|
Same weights, same layers. The only decoder-unique component is a
|
|
self-conditioning MLP.
|
|
|
|
Multimodal support: the model always includes a vision tower (shared with Gemma4).
|
|
Images are encoded through the vision tower and projected into the LM embedding space
|
|
via Gemma4MultimodalEmbedder.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterable, Mapping
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
from transformers import AutoModel
|
|
|
|
from vllm.config import VllmConfig
|
|
from vllm.config.compilation import CUDAGraphMode
|
|
from vllm.distributed.parallel_state import get_tp_group
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.layernorm import RMSNorm
|
|
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
|
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
|
ParallelLMHead,
|
|
)
|
|
from vllm.model_executor.models.gemma4 import Gemma4Model
|
|
from vllm.model_executor.models.gemma4_mm import (
|
|
Gemma4DummyInputsBuilder,
|
|
Gemma4ForConditionalGeneration,
|
|
Gemma4MultimodalEmbedder,
|
|
Gemma4MultiModalProcessor,
|
|
Gemma4ProcessingInfo,
|
|
)
|
|
from vllm.model_executor.models.module_mapping import MultiModelKeys
|
|
from vllm.model_executor.models.transformers.utils import recursive_replace_linear
|
|
from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
|
|
from vllm.multimodal import MULTIMODAL_REGISTRY
|
|
from vllm.platforms import current_platform
|
|
from vllm.v1.outputs import LogprobsTensors
|
|
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
|
|
from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu
|
|
from vllm.v1.worker.gpu.input_batch import InputBatch
|
|
from vllm.v1.worker.gpu.model_states.interface import ModelState
|
|
from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
|
|
from vllm.v1.worker.gpu.sample.output import SamplerOutput
|
|
from vllm.v1.worker.gpu.sample.penalties import use_penalty
|
|
from vllm.v1.worker.gpu.states import RequestState
|
|
|
|
from .interfaces import (
|
|
SupportsMultiModal,
|
|
SupportsPP,
|
|
SupportsQuant,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class DiffusionGemmaSelfConditioning(nn.Module):
|
|
"""Gated MLP that processes soft embeddings from the previous denoising step.
|
|
|
|
Structurally identical to Gemma4MLP but with self_conditioning_size
|
|
and post_norm without learned scale.
|
|
"""
|
|
|
|
def __init__(
|
|
self, hidden_size: int, self_conditioning_size: int, eps: float = 1e-6
|
|
):
|
|
super().__init__()
|
|
self.pre_norm = RMSNorm(hidden_size, eps=eps)
|
|
self.post_norm = RMSNorm(hidden_size, eps=eps, has_weight=False)
|
|
self.gate_proj = nn.Linear(hidden_size, self_conditioning_size, bias=False)
|
|
self.up_proj = nn.Linear(hidden_size, self_conditioning_size, bias=False)
|
|
self.down_proj = nn.Linear(self_conditioning_size, hidden_size, bias=False)
|
|
|
|
def forward(
|
|
self,
|
|
inputs_embeds: torch.Tensor,
|
|
soft_embeds: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
x = self.pre_norm(soft_embeds)
|
|
sc_signal = self.down_proj(
|
|
F.gelu(self.gate_proj(x), approximate="tanh") * self.up_proj(x)
|
|
)
|
|
return self.post_norm(inputs_embeds + sc_signal)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Multimodal processing info (overrides Gemma4 config type check)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class DiffusionGemmaProcessingInfo(Gemma4ProcessingInfo):
|
|
"""Processing info for DiffusionGemma.
|
|
|
|
Overrides ``get_hf_config`` to accept ``DiffusionGemmaConfig``
|
|
(which inherits from ``PretrainedConfig``, not ``Gemma4Config``).
|
|
Supports image and video modalities.
|
|
"""
|
|
|
|
def get_hf_config(self):
|
|
# DiffusionGemmaConfig doesn't inherit from Gemma4Config, so we
|
|
# accept any PretrainedConfig here.
|
|
return self.ctx.get_hf_config()
|
|
|
|
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
|
|
# DiffusionGemma supports image and video inputs.
|
|
return {"image": None, "video": None}
|
|
|
|
def get_mm_max_tokens_per_item(
|
|
self, seq_len: int, mm_counts: Mapping[str, int]
|
|
) -> Mapping[str, int] | None:
|
|
return super().get_mm_max_tokens_per_item(seq_len, mm_counts)
|
|
|
|
|
|
@torch.compile(dynamic=True)
|
|
def _softcap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor:
|
|
# fp32 before tanh for numerical stability (matches HF DiffusionGemma).
|
|
# Compiling fuses the cast/div/tanh/mul into one elementwise kernel over
|
|
# the [num_tokens, vocab] logits instead of four separate passes.
|
|
logits = logits.float()
|
|
return torch.tanh(logits / cap) * cap
|
|
|
|
|
|
@MULTIMODAL_REGISTRY.register_processor(
|
|
Gemma4MultiModalProcessor,
|
|
info=DiffusionGemmaProcessingInfo,
|
|
dummy_inputs=Gemma4DummyInputsBuilder,
|
|
)
|
|
class DiffusionGemmaForConditionalGeneration(
|
|
nn.Module,
|
|
SupportsMultiModal,
|
|
SupportsQuant,
|
|
SupportsPP,
|
|
):
|
|
"""DiffusionGemma for vLLM.
|
|
|
|
Single Gemma4 backbone that switches between encoder and decoder mode.
|
|
The encoder path uses standard Gemma4 layers (causal attention, KV write).
|
|
The decoder path uses the same weights with bidirectional attention and
|
|
KV read-only, plus self-conditioning.
|
|
|
|
Always includes a vision tower (same as Gemma4) for image understanding.
|
|
|
|
In practice, the model's forward() dispatches based on the `mode` kwarg
|
|
set by DiffusionGemmaModelState.prepare_inputs().
|
|
"""
|
|
|
|
hf_to_vllm_mapper = WeightsMapper(
|
|
orig_to_new_prefix={
|
|
"model.decoder.": "model.",
|
|
"model.encoder.language_model.": "model.",
|
|
"model.encoder.vision_tower.": "vision_tower.",
|
|
"model.encoder.embed_vision.": "embed_vision.",
|
|
},
|
|
orig_to_new_substr={
|
|
".experts.": ".moe.experts.",
|
|
},
|
|
)
|
|
|
|
packed_modules_mapping = {
|
|
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
|
"gate_up_proj": ["gate_proj", "up_proj"],
|
|
}
|
|
|
|
@staticmethod
|
|
def get_model_state_cls():
|
|
return DiffusionGemmaModelState
|
|
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
super().__init__()
|
|
config = vllm_config.model_config.hf_config
|
|
text_config = vllm_config.model_config.hf_text_config
|
|
self.config = config
|
|
self.model_dtype = vllm_config.model_config.dtype
|
|
|
|
# DiffusionGemma's full-attention layers have NO v_proj — V is
|
|
# computed from k_proj's output (`value_states = key_states` before
|
|
# k_norm in `DiffusionGemmaDecoderTextAttention.forward`). This is
|
|
# the "k_eq_v" variant in our Gemma4 backbone. The checkpoint has no
|
|
# v_proj weights for full-attention layers; without this flag they
|
|
# would silently load with random V projections.
|
|
text_config.attention_k_eq_v = True
|
|
|
|
# ---- Vision tower ----
|
|
vision_config = getattr(config, "vision_config", None)
|
|
if vision_config is not None:
|
|
quant_config = vllm_config.quant_config
|
|
if quant_config and quant_config.get_name() in [
|
|
"bitsandbytes",
|
|
"torchao",
|
|
"compressed-tensors",
|
|
]:
|
|
tower_quant = quant_config
|
|
else:
|
|
quantizable = (
|
|
vision_config.hidden_size % 64 == 0
|
|
and vision_config.intermediate_size % 64 == 0
|
|
)
|
|
tower_quant = quant_config if quantizable else None
|
|
|
|
with self._mark_tower_model(vllm_config, {"image", "video"}):
|
|
self.vision_tower = AutoModel.from_config(config=vision_config)
|
|
self.embed_vision = Gemma4MultimodalEmbedder(
|
|
vision_config,
|
|
text_config,
|
|
quant_config=tower_quant,
|
|
prefix=maybe_prefix(prefix, "embed_vision"),
|
|
)
|
|
recursive_replace_linear(
|
|
self.vision_tower,
|
|
tower_quant,
|
|
prefix=maybe_prefix(prefix, "vision_tower"),
|
|
)
|
|
else:
|
|
self.vision_tower = None
|
|
self.embed_vision = None
|
|
|
|
# ---- Language backbone (Gemma4Model) ----
|
|
# Use maybe_prefix to ensure correct weight name prefixes for
|
|
# quantization. The quantization config uses hf_to_vllm_mapper to
|
|
# match checkpoint weight names to model parameter names.
|
|
self.model = Gemma4Model(
|
|
vllm_config=vllm_config,
|
|
prefix=maybe_prefix(prefix, "model"),
|
|
)
|
|
|
|
self.lm_head = ParallelLMHead(
|
|
num_embeddings=text_config.vocab_size,
|
|
embedding_dim=text_config.hidden_size,
|
|
)
|
|
|
|
if text_config.tie_word_embeddings:
|
|
self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
|
|
|
|
# HF DiffusionGemma applies the final-logit softcap in fp32, before
|
|
# any other processing. Do it manually in `compute_logits` so the
|
|
# LogitsProcessor only handles the lm_head GEMM.
|
|
self.final_logit_softcapping = getattr(
|
|
text_config, "final_logit_softcapping", None
|
|
)
|
|
self.logits_processor = LogitsProcessor(
|
|
text_config.vocab_size,
|
|
soft_cap=None,
|
|
)
|
|
|
|
sc_size = (
|
|
getattr(config, "self_conditioning_size", None)
|
|
or text_config.intermediate_size
|
|
)
|
|
self.self_conditioning = DiffusionGemmaSelfConditioning(
|
|
hidden_size=text_config.hidden_size,
|
|
self_conditioning_size=sc_size,
|
|
eps=getattr(text_config, "rms_norm_eps", 1e-6),
|
|
)
|
|
|
|
self.make_empty_intermediate_tensors = (
|
|
self.model.make_empty_intermediate_tensors
|
|
)
|
|
|
|
def compute_self_conditioning(
|
|
self,
|
|
inputs_embeds: torch.Tensor,
|
|
probs: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
embed_weight = self.model.embed_tokens.weight
|
|
soft_embeds = torch.matmul(
|
|
probs.to(embed_weight.dtype), embed_weight
|
|
) * self.model.normalizer.to(inputs_embeds.dtype)
|
|
return self.self_conditioning(inputs_embeds, soft_embeds)
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Multimodal: reuse Gemma4's image parsing, processing & embedding
|
|
# ------------------------------------------------------------------ #
|
|
# The vision tower, pooler, embed_vision, and their processing logic
|
|
# are architecturally identical to Gemma4. Delegate to avoid
|
|
# maintaining a duplicate copy.
|
|
|
|
_parse_and_validate_image_input = (
|
|
Gemma4ForConditionalGeneration._parse_and_validate_image_input
|
|
)
|
|
_parse_and_validate_video_input = (
|
|
Gemma4ForConditionalGeneration._parse_and_validate_video_input
|
|
)
|
|
_parse_and_validate_multimodal_inputs = (
|
|
Gemma4ForConditionalGeneration._parse_and_validate_multimodal_inputs
|
|
)
|
|
_encoder_chunk = staticmethod(Gemma4ForConditionalGeneration._encoder_chunk)
|
|
_process_image_input = Gemma4ForConditionalGeneration._process_image_input
|
|
_process_video_input = Gemma4ForConditionalGeneration._process_video_input
|
|
embed_multimodal = Gemma4ForConditionalGeneration.embed_multimodal
|
|
|
|
def get_mm_mapping(self) -> MultiModelKeys:
|
|
"""Get the module prefix mapping for multimodal models."""
|
|
return MultiModelKeys.from_string_field(
|
|
language_model="model",
|
|
connector=["embed_vision"],
|
|
tower_model=["vision_tower"],
|
|
)
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Forward
|
|
# ------------------------------------------------------------------ #
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
intermediate_tensors: Any | None = None,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
**kwargs: Any,
|
|
) -> torch.Tensor:
|
|
if intermediate_tensors is not None:
|
|
inputs_embeds = None
|
|
return self.model(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
**kwargs,
|
|
)
|
|
|
|
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
|
|
logits = self.logits_processor(self.lm_head, hidden_states)
|
|
if logits is not None and self.final_logit_softcapping is not None:
|
|
logits = _softcap_logits(logits, self.final_logit_softcapping)
|
|
return logits
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
|
|
"""Load weights from checkpoint.
|
|
|
|
Checkpoint layout (HF DiffusionGemma):
|
|
model.encoder.vision_tower.* → vision tower
|
|
model.encoder.embed_vision.* → vision embedder
|
|
model.encoder.language_model.layers.* → backbone
|
|
model.decoder.layers.* → backbone (tied)
|
|
model.decoder.embed_tokens.* → embeddings
|
|
model.decoder.self_conditioning.* → self-conditioning MLP
|
|
lm_head.* → LM head (tied)
|
|
|
|
We load encoder weights into our single ``Gemma4Model`` backbone,
|
|
skip duplicate decoder backbone weights, handle vision tower and
|
|
self-conditioning separately.
|
|
"""
|
|
|
|
sc_params = dict(
|
|
(n, p)
|
|
for n, p in self.named_parameters()
|
|
if n.startswith("self_conditioning.")
|
|
)
|
|
|
|
# Collect vision tower + embedder parameters AND buffers for manual
|
|
# loading. The HF vision tower registers std_bias / std_scale as
|
|
# buffers (not parameters) when config.standardize is True, so we
|
|
# must include named_buffers() to avoid "not found in model" warnings.
|
|
vision_params: dict[str, torch.Tensor] = {}
|
|
for n, p in self.named_parameters():
|
|
if n.startswith(("vision_tower.", "embed_vision.")):
|
|
vision_params[n] = p
|
|
for n, b in self.named_buffers():
|
|
if n.startswith(("vision_tower.", "embed_vision.")):
|
|
vision_params[n] = b
|
|
|
|
def _remap_weights():
|
|
# Use full weight names (including suffixes like .weight_scale,
|
|
# .weight_packed) for dedup instead of just the base layer name. Critical
|
|
# for quantized checkpoints where each weight has multiple tensors;
|
|
# tracking only base names skips scales as duplicates.
|
|
seen_weights: set[str] = set()
|
|
for name, weight in weights:
|
|
# Self-conditioning lives under model.decoder.self_conditioning.*
|
|
# in the checkpoint but at self_conditioning.* in our model.
|
|
if "self_conditioning" in name:
|
|
sc_name = name.split("self_conditioning.", 1)[1]
|
|
sc_name = "self_conditioning." + sc_name
|
|
if sc_name in sc_params:
|
|
sc_params[sc_name].data.copy_(weight)
|
|
continue
|
|
|
|
# Vision tower: model.encoder.vision_tower.* → vision_tower.*
|
|
# In HF, the vision tower is a sibling of language_model
|
|
# under the encoder module.
|
|
if name.startswith("model.encoder.vision_tower."):
|
|
vt_name = name[len("model.encoder.") :]
|
|
if vt_name in vision_params:
|
|
vision_params[vt_name].data.copy_(weight)
|
|
else:
|
|
logger.warning(
|
|
"Vision tower weight %s (mapped to %s) not found in model",
|
|
name,
|
|
vt_name,
|
|
)
|
|
continue
|
|
|
|
# Vision embedder: model.encoder.embed_vision.* → embed_vision.*
|
|
if name.startswith("model.encoder.embed_vision."):
|
|
ev_name = name[len("model.encoder.") :]
|
|
if ev_name in vision_params:
|
|
vision_params[ev_name].data.copy_(weight)
|
|
else:
|
|
logger.warning(
|
|
"Embed vision weight %s (mapped to %s) not found in model",
|
|
name,
|
|
ev_name,
|
|
)
|
|
continue
|
|
|
|
# Skip vestigial embed_vision.embedding weights.
|
|
if "embed_vision.embedding." in name:
|
|
continue
|
|
|
|
# Encoder backbone → model.*
|
|
if name.startswith("model.encoder.language_model."):
|
|
name = name.replace("model.encoder.language_model.", "model.")
|
|
# Decoder backbone → model.* (skip exact duplicates)
|
|
elif name.startswith("model.decoder."):
|
|
name = name.replace("model.decoder.", "model.")
|
|
|
|
# Skip only if we've seen the exact same weight name (including scales)
|
|
if name in seen_weights:
|
|
continue
|
|
seen_weights.add(name)
|
|
yield name, weight
|
|
|
|
# Delegate to Gemma4ForCausalLM.load_weights for the backbone,
|
|
# which handles stacked params, MoE, k_eq_v, etc.
|
|
# Temporarily set self.config to text_config since Gemma4's
|
|
# load_weights expects it (e.g. tie_word_embeddings, layer_types).
|
|
from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
|
|
|
|
saved_config = self.config
|
|
self.config = self.model.config
|
|
try:
|
|
Gemma4ForCausalLM.load_weights(self, _remap_weights())
|
|
finally:
|
|
self.config = saved_config
|
|
|
|
@classmethod
|
|
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
|
|
if modality == "image":
|
|
return "<image_soft_token>"
|
|
if modality == "video":
|
|
return "<|video|>"
|
|
raise ValueError(f"Unsupported modality: {modality}")
|
|
|
|
|
|
@torch.compile(dynamic=True)
|
|
def _compute_num_rejected(
|
|
num_logits: torch.Tensor,
|
|
num_sampled: torch.Tensor,
|
|
query_start_loc: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
query_lens = query_start_loc[1:] - query_start_loc[:-1]
|
|
num_rejected = num_logits - num_sampled
|
|
is_denoise = (num_logits > 0) & (num_sampled == 0)
|
|
return torch.where(is_denoise, query_lens, num_rejected)
|
|
|
|
|
|
@torch.compile(dynamic=True)
|
|
def _compiled_sample_step(
|
|
# Logits from the model [num_decode * CL, vocab]
|
|
logits: torch.Tensor,
|
|
# Request mapping
|
|
decode_slots: torch.Tensor, # [num_decode] int64 → slot indices
|
|
decode_idx: torch.Tensor, # [num_decode] int64 → position in num_reqs
|
|
all_slots: torch.Tensor, # [num_reqs] int64 → all slot indices
|
|
valid_canvas_len: torch.Tensor, # [num_decode] int64 → real canvas length (<=CL)
|
|
# State tensors (modified in-place)
|
|
canvas: torch.Tensor, # [max_num_reqs, CL]
|
|
argmax_canvas: torch.Tensor, # [max_num_reqs, CL]
|
|
step_tensor: torch.Tensor, # [max_num_reqs]
|
|
is_encoder_phase: torch.Tensor, # [max_num_reqs]
|
|
confident_tensor: torch.Tensor, # [max_num_reqs]
|
|
sc_embeds: torch.Tensor, # [max_num_reqs, CL, hidden]
|
|
embed_weight: torch.Tensor, # [vocab, hidden]
|
|
normalizer: torch.Tensor,
|
|
history: torch.Tensor, # [max_num_reqs, ST, CL]
|
|
history_len_tensor: torch.Tensor, # [max_num_reqs]
|
|
# Output tensors (modified in-place)
|
|
sampled: torch.Tensor, # [num_reqs, CL]
|
|
num_sampled: torch.Tensor, # [num_reqs]
|
|
draft_tokens: torch.Tensor, # [max_num_reqs, >=CL]
|
|
# Scalar config
|
|
max_denoising_steps: float,
|
|
t_min: float,
|
|
t_max: float,
|
|
confidence_threshold: float,
|
|
vocab_size: int,
|
|
CL: int,
|
|
ST: int,
|
|
# Sampler config
|
|
entropy_bound: float,
|
|
# Tensor-parallel vocab sharding for the self-conditioning matmul.
|
|
# ``embed_weight`` is vocab-sharded ([vocab/tp, hidden]) while ``probs``
|
|
# spans the full vocab; [sc_vocab_start, sc_vocab_end) is this rank's slice.
|
|
sc_vocab_start: int,
|
|
sc_vocab_end: int,
|
|
tp_size: int,
|
|
tp_group_name: str,
|
|
) -> torch.Tensor:
|
|
"""Compiled decode step: temperature → Gumbel sample → probs/confidence →
|
|
accept/renoise → convergence, all as vectorized PyTorch ops.
|
|
|
|
Returns the temperature-scaled logits ``[num_decode, CL, vocab]`` so the
|
|
caller can compute logprobs outside the compiled region."""
|
|
num_decode = decode_slots.shape[0]
|
|
device = decode_slots.device
|
|
|
|
# ---- Phase 1: Temperature schedule ----
|
|
steps_f = step_tensor[decode_slots].float()
|
|
remaining = (max_denoising_steps - steps_f).clamp(min=1.0)
|
|
temp = t_min + (t_max - t_min) * (remaining / max_denoising_steps)
|
|
|
|
# ---- Phase 2: Temperature scaling + Gumbel-max sampling ----
|
|
logits_3d = logits.reshape(num_decode, CL, -1).float()
|
|
scaled = logits_3d / temp[:, None, None].clamp(min=1e-10)
|
|
|
|
# Gumbel-max trick: argmax(logits/T + Gumbel) ~ sample from softmax(logits/T)
|
|
u = torch.rand_like(scaled).clamp(min=1e-20)
|
|
gumbel = -torch.log(-torch.log(u))
|
|
# Zero noise when temp==0 (greedy)
|
|
noisy = scaled + gumbel * (temp[:, None, None] > 0).float()
|
|
new_tokens = noisy.view(-1, noisy.shape[-1]).argmax(dim=-1).view(num_decode, CL)
|
|
argmax_tokens = (
|
|
scaled.view(-1, scaled.shape[-1]).argmax(dim=-1).view(num_decode, CL)
|
|
)
|
|
|
|
# ---- Phase 3: Probs, self-conditioning, confidence ----
|
|
log_probs = scaled.log_softmax(dim=-1)
|
|
probs = log_probs.exp()
|
|
|
|
token_entropy = -(probs * log_probs).sum(dim=-1) # [num_decode, CL]
|
|
# A canvas truncated near max_model_len is zero-padded up to CL by the
|
|
# caller; those padded rows are uniform (max entropy, argmax 0), so they
|
|
# never trigger early convergence and are stable, and only the real
|
|
# ``valid_canvas_len`` tokens are committed (num_sampled below).
|
|
mean_entropy = token_entropy.mean(dim=-1) # [num_decode]
|
|
confident_tensor[decode_slots] = mean_entropy < confidence_threshold
|
|
|
|
# ---- Phase 4: Entropy-bound acceptance mask ----
|
|
sorted_ent, sorted_idx = torch.sort(token_entropy, dim=-1)
|
|
cumsum_ent = torch.cumsum(sorted_ent, dim=-1)
|
|
cummax_ent = torch.cummax(sorted_ent, dim=-1).values
|
|
sorted_mask = (cumsum_ent - cummax_ent) <= entropy_bound
|
|
eb_mask = torch.zeros_like(sorted_mask)
|
|
eb_mask.scatter_(1, sorted_idx, sorted_mask)
|
|
|
|
# ---- Phase 5: Post-sample ----
|
|
is_commit = is_encoder_phase[decode_slots] # [num_decode]
|
|
is_denoise = ~is_commit
|
|
cur_step = step_tensor[decode_slots].float()
|
|
|
|
# Step update: +1 for denoise, reset to 0 for commit
|
|
new_step_val = torch.where(
|
|
is_denoise,
|
|
(cur_step + 1).to(step_tensor.dtype),
|
|
step_tensor.new_zeros(num_decode),
|
|
)
|
|
step_tensor[decode_slots] = new_step_val
|
|
|
|
# Random tokens for renoise / canvas reinit
|
|
random_tokens = torch.randint(
|
|
0, vocab_size, (num_decode, CL), device=device, dtype=canvas.dtype
|
|
)
|
|
|
|
# Compute denoise canvas (accept/renoise)
|
|
denoise_canvas = torch.where(eb_mask, new_tokens, random_tokens)
|
|
|
|
# Canvas: commit → random reinit, denoise → accept/renoise result
|
|
canvas[decode_slots] = torch.where(
|
|
is_commit.unsqueeze(1), random_tokens, denoise_canvas
|
|
)
|
|
|
|
# History: write argmax_tokens for denoise requests at circular position
|
|
hist_len = history_len_tensor[decode_slots]
|
|
write_pos = hist_len % ST
|
|
for i in range(ST):
|
|
write_here = ((write_pos == i) & is_denoise).unsqueeze(1)
|
|
history[decode_slots, i] = torch.where(
|
|
write_here, argmax_tokens, history[decode_slots, i]
|
|
)
|
|
|
|
# Argmax canvas: update for denoise, preserve for commit
|
|
argmax_canvas[decode_slots] = torch.where(
|
|
is_denoise.unsqueeze(1), argmax_tokens, argmax_canvas[decode_slots]
|
|
)
|
|
|
|
# History length: increment for denoise, reset for commit
|
|
new_hist_len = torch.where(is_denoise, hist_len + 1, hist_len.new_zeros(num_decode))
|
|
history_len_tensor[decode_slots] = new_hist_len
|
|
|
|
# Sampled output: commit → emit argmax_canvas, denoise → 0 (pre-zeroed)
|
|
sampled[decode_idx] = argmax_canvas[decode_slots].to(
|
|
sampled.dtype
|
|
) * is_commit.unsqueeze(1).to(sampled.dtype)
|
|
# Commit only the real canvas length (== CL except for a canvas truncated
|
|
# near max_model_len); the padded tail positions are never emitted.
|
|
num_sampled[decode_idx] = is_commit.to(num_sampled.dtype) * valid_canvas_len.to(
|
|
num_sampled.dtype
|
|
)
|
|
|
|
# ---- Phase 6: Stability + convergence ----
|
|
ref = history[decode_slots, 0]
|
|
mismatch = torch.zeros(num_decode, device=device, dtype=torch.int32)
|
|
for h in range(1, ST):
|
|
mismatch = mismatch + (ref != history[decode_slots, h]).sum(dim=-1).int()
|
|
stable = mismatch == 0
|
|
|
|
step_after = step_tensor[decode_slots]
|
|
converged = (stable & confident_tensor[decode_slots] & (new_hist_len >= ST)) | (
|
|
step_after >= max_denoising_steps
|
|
)
|
|
# Commit done → denoise next (False); denoise converged → commit next (True)
|
|
is_encoder_phase[decode_slots] = torch.where(
|
|
is_commit, is_commit.new_zeros(num_decode), converged
|
|
)
|
|
|
|
# SC soft embedding: store ``probs @ embed_weight`` (the value the next step's
|
|
# self-conditioning MLP consumes) only for slots that will denoise next — i.e.
|
|
# this step denoised AND it isn't about to commit (is_encoder_phase now False).
|
|
# Masking here (rather than in the consumer) lets _apply_self_conditioning read
|
|
# sc_embeds directly. Storing the [.., hidden] soft embed instead of the full
|
|
# [.., vocab] probs avoids a giant persistent buffer.
|
|
sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None]
|
|
# Self-conditioning soft embed = probs @ embed_tokens.weight. Under tensor
|
|
# parallelism the embedding is vocab-sharded ([vocab/tp, hidden]) while
|
|
# probs spans the full vocab, so each rank multiplies its local vocab slice
|
|
# [sc_vocab_start, sc_vocab_end) and the partials are summed across ranks.
|
|
local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype)
|
|
soft_embeds = torch.matmul(
|
|
local_probs, embed_weight[: sc_vocab_end - sc_vocab_start]
|
|
)
|
|
if tp_size > 1:
|
|
soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name)
|
|
soft_embeds = soft_embeds * normalizer
|
|
sc_embeds[decode_slots] = soft_embeds * sc_keep
|
|
|
|
# Overwrite canvas with argmax for newly converged denoise requests
|
|
newly_converged = (converged & is_denoise).unsqueeze(1)
|
|
canvas[decode_slots] = torch.where(
|
|
newly_converged, argmax_canvas[decode_slots], canvas[decode_slots]
|
|
)
|
|
|
|
# ---- Phase 7: Copy canvas → draft_tokens for all slots ----
|
|
draft_tokens[all_slots, :CL] = canvas[all_slots]
|
|
|
|
return scaled
|
|
|
|
|
|
class DiffusionGemmaRequestStates:
|
|
"""Pre-allocated GPU tensors for DiffusionGemma per-request state.
|
|
|
|
Follows the indexed-slot pattern used by ``RequestState``.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
max_num_reqs: int,
|
|
canvas_length: int,
|
|
vocab_size: int,
|
|
max_denoising_steps: int,
|
|
device: torch.device,
|
|
hidden_size: int,
|
|
stability_threshold: int,
|
|
):
|
|
self.max_num_reqs = max_num_reqs
|
|
self.canvas_length = canvas_length
|
|
self.vocab_size = vocab_size
|
|
self.max_denoising_steps = max_denoising_steps
|
|
self.stability_threshold = stability_threshold
|
|
self.device = device
|
|
|
|
self.is_encoder_phase = torch.zeros(
|
|
max_num_reqs, dtype=torch.bool, device=device
|
|
)
|
|
# Canvas tokens [max_num_reqs, canvas_length]
|
|
self.canvas = torch.zeros(
|
|
max_num_reqs, canvas_length, dtype=torch.int64, device=device
|
|
)
|
|
# Step counter (counts up from 0 to max_denoising_steps)
|
|
self.step = torch.zeros(
|
|
max_num_reqs,
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
# Accepted canvas history for stability check
|
|
self.accepted_canvas_history = torch.zeros(
|
|
max_num_reqs,
|
|
stability_threshold,
|
|
canvas_length,
|
|
dtype=torch.int64,
|
|
device=device,
|
|
)
|
|
self.accepted_canvas_history_len = torch.zeros(
|
|
max_num_reqs, dtype=torch.int32, device=device
|
|
)
|
|
# Latest argmax(processed_logits) per slot — what we COMMIT.
|
|
# NOT `current_canvas` (which is the post-renoise stochastic input for
|
|
# the next denoise step). We keep this separate from `canvas` because
|
|
# canvas gets renoised in-place during denoise, while argmax_canvas is
|
|
# the deterministic best-guess we ultimately emit.
|
|
self.argmax_canvas = torch.zeros(
|
|
max_num_reqs, canvas_length, dtype=torch.int64, device=device
|
|
)
|
|
|
|
# Per-slot prompt length (set by add_request).
|
|
self.prompt_len = torch.zeros(
|
|
max_num_reqs,
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
|
|
# Per-slot confidence flag, set by the sampler each step.
|
|
self.confident = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
|
|
|
|
# Per-slot self-conditioning soft embedding (probs @ embed_weight) from
|
|
# the previous denoise step. Storing the [.., hidden] soft embed instead
|
|
# of the full [.., vocab] distribution shrinks this buffer by
|
|
# vocab/hidden (~170x) and moves the matmul to denoise time; the result
|
|
# is identical (SC consumes probs @ embed_weight anyway).
|
|
self.self_conditioning_embeds = torch.zeros(
|
|
max_num_reqs, canvas_length, hidden_size, dtype=torch.float32, device=device
|
|
)
|
|
|
|
def init_canvas(self, slot_indices_np: np.ndarray) -> None:
|
|
"""Initialize canvas with random tokens for the given slots."""
|
|
n = slot_indices_np.shape[0]
|
|
self.canvas[slot_indices_np] = torch.randint(
|
|
0,
|
|
self.vocab_size,
|
|
(n, self.canvas_length),
|
|
dtype=torch.int64,
|
|
device=self.device,
|
|
)
|
|
|
|
def add_request(self, slot_idx: int) -> None:
|
|
self.is_encoder_phase[slot_idx] = True
|
|
self.init_canvas(torch.tensor([slot_idx], device=self.device))
|
|
self.step[slot_idx] = 0
|
|
self.accepted_canvas_history_len[slot_idx] = 0
|
|
self.self_conditioning_embeds[slot_idx] = 0
|
|
|
|
def remove_request(self, slot_idx: int) -> None:
|
|
self.is_encoder_phase[slot_idx] = False
|
|
self.accepted_canvas_history_len[slot_idx] = 0
|
|
self.self_conditioning_embeds[slot_idx] = 0
|
|
|
|
|
|
class DiffusionGemmaModelState(ModelState):
|
|
"""ModelState for DiffusionGemma.
|
|
|
|
Single Gemma4 backbone in two modes:
|
|
- encoder mode (num_draft_tokens == 0): causal attention, writes KV
|
|
- decoder mode (num_draft_tokens > 0): bidirectional attention, reads KV
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
model: nn.Module,
|
|
encoder_cache: Any,
|
|
device: torch.device,
|
|
) -> None:
|
|
super().__init__(vllm_config, model, encoder_cache, device)
|
|
|
|
# Per-step MM data produced by get_mm_embeddings and consumed by
|
|
# prepare_inputs. Stored as raw (mm_embeds, is_mm_embed) so that
|
|
# prepare_inputs can call embed_input_ids directly into the
|
|
# persistent _inputs_embeds_buf, avoiding the intermediate copy
|
|
# through encoder_runner.inputs_embeds.
|
|
self._pending_mm_embeds: tuple[list[torch.Tensor], torch.Tensor] | None = None
|
|
|
|
diffusion_config = vllm_config.diffusion_config
|
|
canvas_length = diffusion_config.canvas_length if diffusion_config else 32
|
|
|
|
text_config = self.model_config.hf_text_config
|
|
self.gen_config = self.model_config.try_get_generation_config()
|
|
max_denoising_steps = (
|
|
diffusion_config.max_denoising_steps if diffusion_config else None
|
|
) or self.gen_config.get("max_denoising_steps", 48)
|
|
self.diffusion_states = DiffusionGemmaRequestStates(
|
|
max_num_reqs=self.max_num_reqs,
|
|
canvas_length=canvas_length,
|
|
vocab_size=self.model_config.get_vocab_size(),
|
|
max_denoising_steps=max_denoising_steps,
|
|
device=device,
|
|
hidden_size=text_config.hidden_size,
|
|
# In Transformers, `stability_threshold=1` (the default) means the current
|
|
# step must match the previous step. In vLLM, the history buffer includes
|
|
# the current step, so we add 1 to match the same behavior.
|
|
stability_threshold=self.gen_config["stability_threshold"] + 1,
|
|
)
|
|
self._req_id_to_index: dict[str, int] = {}
|
|
|
|
# Persistent buffer for per-request causal flags, updated in-place
|
|
# so FULL CUDA graph replay sees the latest values.
|
|
self._causal_buf = torch.zeros(
|
|
self.max_num_reqs, dtype=torch.bool, device=device
|
|
)
|
|
|
|
# Persistent inputs_embeds buffer — required so FULL CUDA graph
|
|
# capture and runtime point at the SAME memory address.
|
|
# `prepare_dummy_inputs` (capture path) and `prepare_inputs` (runtime
|
|
# path) both must hand the captured graph a tensor at this address.
|
|
self._inputs_embeds_buf = torch.zeros(
|
|
self.max_num_tokens,
|
|
text_config.hidden_size,
|
|
dtype=self.model_config.dtype,
|
|
device=device,
|
|
)
|
|
|
|
def get_supported_generation_tasks(self):
|
|
return ("generate",)
|
|
|
|
def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None:
|
|
diffusion_config = self.vllm_config.diffusion_config
|
|
gen = self.gen_config
|
|
sampler_cfg = gen.get("sampler_config") or {}
|
|
if "EntropyBound" not in sampler_cfg.get("_cls_name", ""):
|
|
raise ValueError("DiffusionGemma requires an EntropyBound sampler_config")
|
|
entropy_bound = sampler_cfg.get("entropy_bound")
|
|
if entropy_bound is None or entropy_bound <= 0:
|
|
raise ValueError(
|
|
f"entropy_bound must be a positive float (got {entropy_bound})"
|
|
)
|
|
# The self-conditioning matmul (probs @ embed_tokens.weight) runs over a
|
|
# vocab-parallel embedding shard. Hand the sampler this rank's vocab
|
|
# slice and TP group so it can all-reduce the partial products.
|
|
embed_tokens = self.model.model.embed_tokens
|
|
shard = embed_tokens.shard_indices
|
|
tp_group = get_tp_group()
|
|
return DiffusionSampler(
|
|
sampler=sampler,
|
|
diffusion_config=diffusion_config,
|
|
vocab_size=self.model_config.get_vocab_size(),
|
|
diffusion_states=self.diffusion_states,
|
|
t_min=gen["t_min"],
|
|
t_max=gen["t_max"],
|
|
entropy_bound=entropy_bound,
|
|
confidence_threshold=gen["confidence_threshold"],
|
|
embed_weight=embed_tokens.weight,
|
|
normalizer=self.model.model.normalizer,
|
|
sc_vocab_start=shard.org_vocab_start_index,
|
|
sc_vocab_end=shard.org_vocab_end_index,
|
|
tp_size=tp_group.world_size,
|
|
tp_group_name=tp_group.unique_name,
|
|
), None
|
|
|
|
def apply_staged_writes(self) -> None:
|
|
pass
|
|
|
|
def add_request(self, req_index: int, new_req_data: Any) -> None:
|
|
self._req_id_to_index[new_req_data.req_id] = req_index
|
|
self.diffusion_states.add_request(req_index)
|
|
if not new_req_data.req_id.startswith("_warmup_"):
|
|
prompt_len = len(new_req_data.prompt_token_ids)
|
|
self.diffusion_states.prompt_len[req_index] = prompt_len
|
|
|
|
def remove_request(self, req_id: str) -> None:
|
|
idx = self._req_id_to_index.pop(req_id, None)
|
|
if idx is not None:
|
|
self.diffusion_states.remove_request(idx)
|
|
|
|
def get_mm_embeddings(
|
|
self,
|
|
scheduled_encoder_inputs: dict[str, list[int]],
|
|
input_batch: InputBatch,
|
|
req_states: RequestState,
|
|
) -> torch.Tensor | None:
|
|
if not self.supports_mm_inputs:
|
|
return None
|
|
|
|
mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
|
|
scheduled_encoder_inputs
|
|
)
|
|
if mm_kwargs:
|
|
encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
|
|
self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))
|
|
|
|
mm_embeds, is_mm_embed = self.gather_mm_embeddings(input_batch)
|
|
|
|
if not mm_embeds:
|
|
# No MM tokens in this batch (e.g. all-decode step).
|
|
# prepare_inputs will use embed_input_ids (text-only) directly.
|
|
self._pending_mm_embeds = None
|
|
return None
|
|
|
|
# Stash raw MM ingredients for prepare_inputs to merge directly
|
|
# into the persistent buffer, avoiding the intermediate copy
|
|
# through encoder_runner.inputs_embeds.
|
|
self._pending_mm_embeds = (mm_embeds, is_mm_embed)
|
|
return None
|
|
|
|
def _apply_self_conditioning(
|
|
self,
|
|
decode_slots_np: np.ndarray,
|
|
decode_idx_np: np.ndarray,
|
|
query_start_loc_np: np.ndarray,
|
|
inputs_embeds: torch.Tensor,
|
|
sc_embeds: torch.Tensor,
|
|
) -> None:
|
|
# One self-conditioning MLP call per decode request, over that request's
|
|
# query span [start, end) = its canvas. The span is the full canvas (CL)
|
|
# or, for the final canvas truncated near max_model_len, fewer than CL
|
|
# positions. sc_embeds already holds probs @ embed_weight from the prior
|
|
# denoise step, masked to zero by the sampler for slots not denoising
|
|
# this step; only the MLP runs here. CPU metadata -> no GPU syncs.
|
|
for slot, idx in zip(decode_slots_np.tolist(), decode_idx_np.tolist()):
|
|
start = int(query_start_loc_np[idx])
|
|
end = int(query_start_loc_np[idx + 1])
|
|
canvas = slice(start, end)
|
|
soft = sc_embeds[slot, : end - start]
|
|
inputs_embeds[canvas] = self.model.self_conditioning(
|
|
inputs_embeds[canvas], soft.to(inputs_embeds.dtype)
|
|
)
|
|
|
|
def prepare_inputs(self, input_batch, req_states) -> dict[str, Any]:
|
|
states = self.diffusion_states
|
|
num_tokens = input_batch.num_tokens
|
|
num_reqs = input_batch.num_reqs
|
|
|
|
# Write into the PERSISTENT inputs_embeds buffer so FULL CUDA graph
|
|
# replay sees the latest values at the captured address.
|
|
num_tokens_padded = input_batch.num_tokens_after_padding
|
|
inputs_embeds = self._inputs_embeds_buf[:num_tokens_padded]
|
|
|
|
# Populate embeddings: merge MM features when available,
|
|
# otherwise embed input_ids as text-only.
|
|
input_ids = input_batch.input_ids[:num_tokens]
|
|
if self._pending_mm_embeds is not None:
|
|
mm_embeds, is_mm_embed = self._pending_mm_embeds
|
|
self._pending_mm_embeds = None
|
|
inputs_embeds[:num_tokens].copy_(
|
|
self.model.embed_input_ids(
|
|
input_ids,
|
|
multimodal_embeddings=mm_embeds,
|
|
is_multimodal=is_mm_embed,
|
|
)
|
|
)
|
|
else:
|
|
inputs_embeds[:num_tokens].copy_(self.model.embed_input_ids(input_ids))
|
|
|
|
# Apply self-conditioning ONLY for denoising decode requests.
|
|
if input_batch.num_draft_tokens > 0 and self._req_id_to_index:
|
|
slots_np = input_batch.idx_mapping_np[:num_reqs]
|
|
num_logits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1])
|
|
is_decode_indices_np = np.where(num_logits_np > 0)[0]
|
|
self._apply_self_conditioning(
|
|
slots_np[is_decode_indices_np],
|
|
is_decode_indices_np,
|
|
input_batch.query_start_loc_np,
|
|
inputs_embeds,
|
|
states.self_conditioning_embeds,
|
|
)
|
|
|
|
return {"inputs_embeds": inputs_embeds}
|
|
|
|
def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
|
|
# CUDA graph capture path — return a slice of the SAME persistent
|
|
# inputs_embeds buffer that `prepare_inputs` writes to at runtime,
|
|
# so the captured graph and runtime point to identical addresses.
|
|
return {"inputs_embeds": self._inputs_embeds_buf[:num_tokens]}
|
|
|
|
def postprocess_state(
|
|
self, idx_mapping, num_sampled, num_computed_tokens=None
|
|
) -> None:
|
|
return None
|
|
|
|
def prepare_attn(
|
|
self,
|
|
input_batch,
|
|
cudagraph_mode,
|
|
block_tables,
|
|
slot_mappings,
|
|
attn_groups,
|
|
kv_cache_config,
|
|
for_capture=False,
|
|
) -> dict[str, Any]:
|
|
if cudagraph_mode == CUDAGraphMode.FULL:
|
|
num_reqs = input_batch.num_reqs_after_padding
|
|
num_tokens = input_batch.num_tokens_after_padding
|
|
else:
|
|
num_reqs = input_batch.num_reqs
|
|
num_tokens = input_batch.num_tokens
|
|
|
|
query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
|
|
max_query_len = input_batch.num_scheduled_tokens.max().item()
|
|
|
|
# Per-request causal mode: encoder (commit) = causal,
|
|
# denoise = bidirectional. Pass GPU tensor so the attention
|
|
# backend can handle mixed batches.
|
|
actual_num_reqs = input_batch.num_reqs
|
|
slots = input_batch.idx_mapping[:actual_num_reqs]
|
|
# Invariant: the sampler flips is_encoder_phase to False only after a
|
|
# request's FINAL prompt chunk, so a prompt spanning multiple chunks
|
|
# (longer than the token budget) stays causal for every chunk.
|
|
self._causal_buf[:actual_num_reqs] = self.diffusion_states.is_encoder_phase[
|
|
slots
|
|
]
|
|
if actual_num_reqs < num_reqs:
|
|
self._causal_buf[actual_num_reqs:num_reqs] = False
|
|
causal: bool | torch.Tensor = self._causal_buf[:num_reqs]
|
|
|
|
return build_attn_metadata(
|
|
attn_groups=attn_groups,
|
|
num_reqs=num_reqs,
|
|
num_tokens=num_tokens,
|
|
query_start_loc_gpu=input_batch.query_start_loc,
|
|
query_start_loc_cpu=query_start_loc_cpu,
|
|
max_query_len=max_query_len,
|
|
seq_lens=input_batch.seq_lens,
|
|
max_seq_len=self.max_model_len,
|
|
block_tables=block_tables,
|
|
slot_mappings=slot_mappings,
|
|
kv_cache_config=kv_cache_config,
|
|
causal=causal,
|
|
)
|
|
|
|
num_new_sampled_tokens_per_step: int = 0
|
|
|
|
|
|
# Penalty stub for the diffusion path: the runner reads
|
|
# penalties_state.output_bin_counts, and post_update treats None as
|
|
# "no penalty bookkeeping".
|
|
_NO_PENALTIES_STATE = SimpleNamespace(output_bin_counts=None)
|
|
|
|
|
|
class DiffusionSampler:
|
|
"""Batched accept/renoise sampler for DiffusionGemma.
|
|
|
|
Follows the same structure as ``vllm.v1.worker.gpu.sample.sampler.Sampler``:
|
|
decomposed into named methods, all GPU state in pre-allocated buffers,
|
|
no GPU→CPU syncs on the hot path.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
sampler: Any,
|
|
diffusion_config: Any,
|
|
vocab_size: int,
|
|
diffusion_states: DiffusionGemmaRequestStates | None = None,
|
|
*,
|
|
confidence_threshold: float,
|
|
t_min: float,
|
|
t_max: float,
|
|
entropy_bound: float,
|
|
embed_weight: torch.Tensor,
|
|
normalizer: torch.Tensor,
|
|
sc_vocab_start: int = 0,
|
|
sc_vocab_end: int | None = None,
|
|
tp_size: int = 1,
|
|
tp_group_name: str = "",
|
|
):
|
|
self.sampling_states = sampler.sampling_states
|
|
self.req_states = sampler.req_states
|
|
# Self-conditioning soft embed = probs @ embed_weight * normalizer,
|
|
# computed in the sampler (see _compiled_sample_step). ``embed_weight``
|
|
# is the vocab-parallel shard; [sc_vocab_start, sc_vocab_end) is this
|
|
# rank's slice of the full vocab and tp_* drive the cross-rank
|
|
# all-reduce.
|
|
self.embed_weight = embed_weight
|
|
self.normalizer = normalizer
|
|
self.sc_vocab_start = sc_vocab_start
|
|
self.sc_vocab_end = sc_vocab_end if sc_vocab_end is not None else vocab_size
|
|
self.tp_size = tp_size
|
|
self.tp_group_name = tp_group_name
|
|
self.canvas_length = (
|
|
diffusion_config.canvas_length if diffusion_config is not None else 32
|
|
)
|
|
self.t_min = t_min
|
|
self.t_max = t_max
|
|
self.confidence_threshold = confidence_threshold
|
|
self.vocab_size = vocab_size
|
|
self.diffusion_states = diffusion_states
|
|
self.entropy_bound = entropy_bound
|
|
|
|
max_num_reqs = diffusion_states.max_num_reqs
|
|
device = diffusion_states.device
|
|
self._sampled = torch.zeros(
|
|
max_num_reqs,
|
|
self.canvas_length,
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
self._num_sampled = torch.zeros(
|
|
max_num_reqs,
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
self._decode_slots = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
|
|
self._decode_idx = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
|
|
self._query_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
|
|
self._num_logits = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
|
|
|
|
# Per-slot stash for logprobs computed on the converging denoise step.
|
|
# Populated after the post-sample kernel detects convergence; consumed
|
|
# on the subsequent commit step when num_sampled=CANVAS_LEN.
|
|
self._pending_logprobs: dict[int, LogprobsTensors] = {}
|
|
|
|
def add_request(self, req_idx: int, prompt_len: int, sampling_params: Any) -> None:
|
|
if use_penalty(sampling_params):
|
|
logger.warning_once(
|
|
"DiffusionGemma does not support repetition/frequency/presence "
|
|
"penalties; ignoring them for this request."
|
|
)
|
|
# Purge any stale logprobs stashed under this slot by a prior request
|
|
# that was aborted between its converging denoise and commit steps.
|
|
self._pending_logprobs.pop(req_idx, None)
|
|
self.sampling_states.add_request(req_idx, sampling_params)
|
|
|
|
def apply_staged_writes(self) -> None:
|
|
self.sampling_states.apply_staged_writes()
|
|
|
|
@property
|
|
def penalties_state(self):
|
|
# Diffusion applies no penalties. The runner reads
|
|
# penalties_state.output_bin_counts, so expose a stub holding None;
|
|
# post_update treats None bin counts as "no penalty bookkeeping".
|
|
return _NO_PENALTIES_STATE
|
|
|
|
# ------------------------------------------------------------------
|
|
# Prefill
|
|
# ------------------------------------------------------------------
|
|
|
|
def _finish_prefills(
|
|
self, input_batch: Any, prefill_indices_np: np.ndarray
|
|
) -> None:
|
|
"""Transition requests whose prompt completes this step to denoising.
|
|
|
|
Initializes their canvas, seeds draft tokens, and flips
|
|
is_encoder_phase to False. Mid-chunk requests (prompt longer than the
|
|
token budget) are left untouched so is_encoder_phase stays True and
|
|
prepare_attn keeps causal attention for their remaining chunks.
|
|
"""
|
|
states = self.diffusion_states
|
|
done_prefill_np = (
|
|
input_batch.num_computed_prefill_tokens_np[prefill_indices_np]
|
|
+ input_batch.num_scheduled_tokens[prefill_indices_np]
|
|
>= input_batch.prefill_len_np[prefill_indices_np]
|
|
)
|
|
ps = input_batch.idx_mapping_np[prefill_indices_np[done_prefill_np]]
|
|
if len(ps) == 0:
|
|
return
|
|
states.init_canvas(ps)
|
|
self.req_states.draft_tokens[ps, : self.canvas_length] = states.canvas[ps]
|
|
ps_gpu = async_copy_to_gpu(
|
|
ps.astype(np.int64), device=states.is_encoder_phase.device
|
|
)
|
|
states.is_encoder_phase.index_fill_(0, ps_gpu, False)
|
|
|
|
def _handle_prefill(
|
|
self,
|
|
input_batch: Any,
|
|
device: torch.device,
|
|
) -> SamplerOutput:
|
|
num_reqs = input_batch.num_reqs
|
|
self._finish_prefills(input_batch, np.arange(num_reqs))
|
|
sampled = self._sampled[:num_reqs, :1]
|
|
sampled.zero_()
|
|
num_sampled = self._num_sampled[:num_reqs]
|
|
num_sampled.zero_()
|
|
return SamplerOutput(
|
|
sampled_token_ids=sampled,
|
|
logprobs_tensors=None,
|
|
num_nans=None,
|
|
num_sampled=num_sampled,
|
|
num_rejected=num_sampled,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Decode helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _build_output(
|
|
self,
|
|
input_batch: Any,
|
|
sampled: torch.Tensor,
|
|
num_sampled: torch.Tensor,
|
|
per_req_nlogits_np: np.ndarray,
|
|
device: torch.device,
|
|
logprobs_tensors: LogprobsTensors | None = None,
|
|
) -> SamplerOutput:
|
|
"""Compute num_rejected and build SamplerOutput."""
|
|
num_reqs = input_batch.num_reqs
|
|
|
|
self._query_lens.np[:num_reqs] = np.diff(
|
|
input_batch.query_start_loc_np[: num_reqs + 1]
|
|
)
|
|
self._num_logits.np[:num_reqs] = per_req_nlogits_np
|
|
self._query_lens.copy_to_uva()
|
|
self._num_logits.copy_to_uva()
|
|
|
|
num_rejected = _compute_num_rejected(
|
|
self._num_logits.gpu[:num_reqs],
|
|
num_sampled,
|
|
input_batch.query_start_loc[: num_reqs + 1],
|
|
)
|
|
|
|
return SamplerOutput(
|
|
sampled_token_ids=sampled,
|
|
logprobs_tensors=logprobs_tensors,
|
|
num_nans=None,
|
|
num_sampled=num_sampled,
|
|
num_rejected=num_rejected,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Main entry point
|
|
# ------------------------------------------------------------------
|
|
|
|
def __call__(
|
|
self,
|
|
logits: torch.Tensor,
|
|
input_batch: Any,
|
|
draft_logits: torch.Tensor | None = None,
|
|
) -> SamplerOutput:
|
|
num_reqs = input_batch.num_reqs
|
|
device = logits.device
|
|
|
|
if input_batch.num_draft_tokens == 0:
|
|
return self._handle_prefill(input_batch, device)
|
|
|
|
# --- CPU/NumPy setup (outside compile): split decode vs prefill, init
|
|
# canvas for any new prefills, and stage decode slot indices to GPU. ---
|
|
states = self.diffusion_states
|
|
CL = self.canvas_length
|
|
slots_np = input_batch.idx_mapping_np[:num_reqs]
|
|
per_req_nlogits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1])
|
|
|
|
decode_indices_np = np.where(per_req_nlogits_np > 0)[0]
|
|
prefill_indices_np = np.where(per_req_nlogits_np == 0)[0]
|
|
decode_slots_np = slots_np[decode_indices_np]
|
|
|
|
if len(prefill_indices_np) > 0:
|
|
self._finish_prefills(input_batch, prefill_indices_np)
|
|
|
|
num_decode = len(decode_indices_np)
|
|
self._decode_slots.np[:num_decode] = decode_slots_np
|
|
self._decode_idx.np[:num_decode] = decode_indices_np
|
|
self._decode_slots.copy_to_uva()
|
|
self._decode_idx.copy_to_uva()
|
|
decode_slots = self._decode_slots.gpu[:num_decode]
|
|
decode_idx = self._decode_idx.gpu[:num_decode]
|
|
|
|
# Real canvas length per decode request. Equals CL except when a canvas
|
|
# was truncated near max_model_len, in which case the scheduler gave us
|
|
# fewer than CL logits for that request.
|
|
valid_canvas_len_np = per_req_nlogits_np[per_req_nlogits_np > 0]
|
|
valid_canvas_len = async_copy_to_gpu(
|
|
valid_canvas_len_np.astype(np.int64), device=device
|
|
)
|
|
|
|
# Pad any truncated canvas back to CL so the uniform-CL sampler math
|
|
# holds. Phantom (padded) positions are zeroed → uniform logits → high
|
|
# entropy (no premature convergence) and argmax 0 (stable); they are
|
|
# never committed (num_sampled == real length).
|
|
if num_decode > 0 and valid_canvas_len_np.min() < CL:
|
|
ar = torch.arange(CL, device=device)
|
|
starts = valid_canvas_len.cumsum(0) - valid_canvas_len # row offset per req
|
|
valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1) # [num_decode, CL]
|
|
src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1)
|
|
logits = logits[src.reshape(-1)] * valid.reshape(-1, 1).to(logits.dtype)
|
|
|
|
# Clear once: the tiled loop below only scatters its own decode slots,
|
|
# so it must not re-clear earlier tiles' writes.
|
|
sampled = self._sampled[:num_reqs]
|
|
num_sampled = self._num_sampled[:num_reqs]
|
|
sampled.zero_()
|
|
num_sampled.zero_()
|
|
|
|
all_slots = input_batch.idx_mapping[:num_reqs]
|
|
|
|
# Snapshot which slots are committing BEFORE the compiled step runs,
|
|
# since it mutates is_encoder_phase (commit→False, converge→True).
|
|
is_committing = states.is_encoder_phase[decode_slots].clone()
|
|
|
|
slots_np = input_batch.idx_mapping_np[:num_reqs]
|
|
is_decode_np = per_req_nlogits_np > 0
|
|
max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np)
|
|
|
|
# Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in
|
|
# _compiled_sample_step keeps several live [group * CL, vocab] copies, so
|
|
# size each tile to a fraction of free memory to bound the transient at
|
|
# high concurrency. Tiling is bit-identical to a single pass.
|
|
group = max(num_decode, 1)
|
|
if num_decode > 0:
|
|
free, _ = current_platform.mem_get_info()
|
|
# ~10 transient fp32 copies of [group * CL, vocab] inside the step
|
|
# (eager peaks at ~8; pad for allocator overhead and small tensors).
|
|
bytes_per_req = CL * self.vocab_size * 4 * 10
|
|
budget = int(free * 0.5) // max(bytes_per_req, 1)
|
|
group = max(1, min(num_decode, budget))
|
|
|
|
for start_req in range(0, num_decode, group):
|
|
end_req = min(start_req + group, num_decode)
|
|
tile = slice(start_req, end_req)
|
|
tile_slots = decode_slots[tile]
|
|
|
|
scaled = _compiled_sample_step(
|
|
logits[start_req * CL : end_req * CL],
|
|
tile_slots,
|
|
decode_idx[tile],
|
|
all_slots,
|
|
valid_canvas_len[tile],
|
|
# State
|
|
states.canvas,
|
|
states.argmax_canvas,
|
|
states.step,
|
|
states.is_encoder_phase,
|
|
states.confident,
|
|
states.self_conditioning_embeds,
|
|
self.embed_weight,
|
|
self.normalizer,
|
|
states.accepted_canvas_history,
|
|
states.accepted_canvas_history_len,
|
|
# Output
|
|
sampled,
|
|
num_sampled,
|
|
self.req_states.draft_tokens,
|
|
# Config
|
|
max_denoising_steps=float(states.max_denoising_steps),
|
|
t_min=self.t_min,
|
|
t_max=self.t_max,
|
|
confidence_threshold=self.confidence_threshold,
|
|
vocab_size=self.vocab_size,
|
|
CL=CL,
|
|
ST=states.stability_threshold,
|
|
entropy_bound=self.entropy_bound,
|
|
sc_vocab_start=self.sc_vocab_start,
|
|
sc_vocab_end=self.sc_vocab_end,
|
|
tp_size=self.tp_size,
|
|
tp_group_name=self.tp_group_name,
|
|
)
|
|
|
|
# Logprobs for denoise steps that just converged (is_encoder_phase
|
|
# flipped False→True), stashed per tile so `scaled` is freed each tile.
|
|
if max_num_logprobs >= 0:
|
|
converged_mask = states.is_encoder_phase[tile_slots]
|
|
just_converged = converged_mask & ~is_committing[tile]
|
|
if just_converged.any():
|
|
flat_logits = scaled.reshape(-1, scaled.shape[-1])
|
|
argmax_tokens = scaled.argmax(dim=-1)
|
|
for local_idx in just_converged.nonzero(as_tuple=True)[0]:
|
|
li = local_idx.item()
|
|
slot = tile_slots[local_idx]
|
|
# Stash only the real canvas positions (== CL unless this
|
|
# canvas was truncated near max_model_len); padded tail
|
|
# positions are never emitted.
|
|
k_i = int(valid_canvas_len_np[start_req + li])
|
|
pos = li * CL
|
|
self._pending_logprobs[slot.item()] = compute_topk_logprobs(
|
|
flat_logits[pos : pos + k_i],
|
|
max_num_logprobs,
|
|
argmax_tokens[local_idx][:k_i],
|
|
)
|
|
|
|
# Commit steps: is_committing was True at entry. Reassemble previously
|
|
# stashed logprobs and attach to SamplerOutput.
|
|
logprobs_tensors = None
|
|
if max_num_logprobs >= 0 and is_committing.any() and self._pending_logprobs:
|
|
parts_ids, parts_lp, parts_ranks = [], [], []
|
|
cu_gen: list[int] = []
|
|
flat_offset = 0
|
|
for i in range(num_reqs):
|
|
cu_gen.append(flat_offset)
|
|
slot = int(slots_np[i])
|
|
if is_decode_np[i] and slot in self._pending_logprobs:
|
|
lp = self._pending_logprobs.pop(slot)
|
|
parts_ids.append(lp.logprob_token_ids)
|
|
parts_lp.append(lp.logprobs)
|
|
parts_ranks.append(lp.selected_token_ranks)
|
|
flat_offset += lp.logprobs.shape[0]
|
|
if parts_ids:
|
|
logprobs_tensors = LogprobsTensors(
|
|
logprob_token_ids=torch.cat(parts_ids),
|
|
logprobs=torch.cat(parts_lp),
|
|
selected_token_ranks=torch.cat(parts_ranks),
|
|
cu_num_generated_tokens=cu_gen,
|
|
)
|
|
|
|
return self._build_output(
|
|
input_batch,
|
|
sampled,
|
|
num_sampled,
|
|
per_req_nlogits_np,
|
|
device,
|
|
logprobs_tensors=logprobs_tensors,
|
|
)
|