| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DiffusionGemma model, ModelState, and Sampler for vLLM. |
| 4 | |
| 5 | Single Gemma4 backbone run in two modes (like YOCO): |
| 6 | - encoder mode: causal attention, writes KV cache |
| 7 | - decoder mode: bidirectional attention, reads encoder KV, doesn't write |
| 8 | |
| 9 | Same weights, same layers. The only decoder-unique component is a |
| 10 | self-conditioning MLP. |
| 11 | |
| 12 | Multimodal support: the model always includes a vision tower (shared with Gemma4). |
| 13 | Images are encoded through the vision tower and projected into the LM embedding space |
| 14 | via Gemma4MultimodalEmbedder. |
| 15 | """ |
| 16 | |
| 17 | from __future__ import annotations |
| 18 | |
| 19 | from collections.abc import Iterable, Mapping, Sequence |
| 20 | from types import SimpleNamespace |
| 21 | from typing import Any |
| 22 | |
| 23 | import numpy as np |
| 24 | import torch |
| 25 | import torch._dynamo |
| 26 | from torch import nn |
| 27 | from torch.nn import functional as F |
| 28 | from transformers import AutoModel |
| 29 | |
| 30 | from vllm.config import VllmConfig |
| 31 | from vllm.config.compilation import CUDAGraphMode |
| 32 | from vllm.distributed.parallel_state import get_tp_group |
| 33 | from vllm.logger import init_logger |
| 34 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 35 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 36 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 37 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 38 | ParallelLMHead, |
| 39 | ) |
| 40 | from vllm.model_executor.models.diffusion_gemma_sampler import ( |
| 41 | sample_row_stats, |
| 42 | sample_row_stats_reference, |
| 43 | ) |
| 44 | from vllm.model_executor.models.gemma4 import Gemma4Model |
| 45 | from vllm.model_executor.models.gemma4_mm import ( |
| 46 | Gemma4DummyInputsBuilder, |
| 47 | Gemma4ForConditionalGeneration, |
| 48 | Gemma4MultimodalEmbedder, |
| 49 | Gemma4MultiModalProcessor, |
| 50 | Gemma4ProcessingInfo, |
| 51 | ) |
| 52 | from vllm.model_executor.models.module_mapping import MultiModelKeys |
| 53 | from vllm.model_executor.models.transformers.utils import recursive_replace_linear |
| 54 | from vllm.model_executor.models.utils import ( |
| 55 | AutoWeightsLoader, |
| 56 | WeightsMapper, |
| 57 | maybe_prefix, |
| 58 | ) |
| 59 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 60 | from vllm.platforms import current_platform |
| 61 | from vllm.utils.torch_utils import async_tensor_h2d |
| 62 | from vllm.v1.outputs import LogprobsTensors |
| 63 | from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p |
| 64 | from vllm.v1.worker.gpu.attn_utils import build_attn_metadata |
| 65 | from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor |
| 66 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 67 | from vllm.v1.worker.gpu.model_states.interface import ModelState |
| 68 | from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores |
| 69 | from vllm.v1.worker.gpu.sample.output import SamplerOutput |
| 70 | from vllm.v1.worker.gpu.sample.penalties import use_penalty |
| 71 | from vllm.v1.worker.gpu.states import RequestState |
| 72 | |
| 73 | from .interfaces import ( |
| 74 | SupportsMultiModal, |
| 75 | SupportsQuant, |
| 76 | ) |
| 77 | |
| 78 | logger = init_logger(__name__) |
| 79 | |
| 80 | |
| 81 | class DiffusionGemmaSelfConditioning(nn.Module): |
| 82 | """Gated MLP that processes soft embeddings from the previous denoising step. |
| 83 | |
| 84 | Structurally identical to Gemma4MLP but with self_conditioning_size |
| 85 | and post_norm without learned scale. |
| 86 | """ |
| 87 | |
| 88 | def __init__( |
| 89 | self, hidden_size: int, self_conditioning_size: int, eps: float = 1e-6 |
| 90 | ): |
| 91 | super().__init__() |
| 92 | self.pre_norm = RMSNorm(hidden_size, eps=eps) |
| 93 | self.post_norm = RMSNorm(hidden_size, eps=eps, has_weight=False) |
| 94 | self.gate_proj = nn.Linear(hidden_size, self_conditioning_size, bias=False) |
| 95 | self.up_proj = nn.Linear(hidden_size, self_conditioning_size, bias=False) |
| 96 | self.down_proj = nn.Linear(self_conditioning_size, hidden_size, bias=False) |
| 97 | |
| 98 | def forward( |
| 99 | self, |
| 100 | inputs_embeds: torch.Tensor, |
| 101 | soft_embeds: torch.Tensor, |
| 102 | ) -> torch.Tensor: |
| 103 | x = self.pre_norm(soft_embeds) |
| 104 | sc_signal = self.down_proj( |
| 105 | F.gelu(self.gate_proj(x), approximate="tanh") * self.up_proj(x) |
| 106 | ) |
| 107 | return self.post_norm(inputs_embeds + sc_signal) |
| 108 | |
| 109 | |
| 110 | # --------------------------------------------------------------------------- |
| 111 | # Multimodal processing info (overrides Gemma4 config type check) |
| 112 | # --------------------------------------------------------------------------- |
| 113 | |
| 114 | |
| 115 | class DiffusionGemmaProcessingInfo(Gemma4ProcessingInfo): |
| 116 | """Processing info for DiffusionGemma. |
| 117 | |
| 118 | Overrides ``get_hf_config`` to accept ``DiffusionGemmaConfig`` |
| 119 | (which inherits from ``PreTrainedConfig``, not ``Gemma4Config``). |
| 120 | Supports image and video modalities. |
| 121 | """ |
| 122 | |
| 123 | def get_hf_config(self): |
| 124 | # DiffusionGemmaConfig doesn't inherit from Gemma4Config, so we |
| 125 | # accept any PreTrainedConfig here. |
| 126 | return self.ctx.get_hf_config() |
| 127 | |
| 128 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 129 | # DiffusionGemma supports image and video inputs. |
| 130 | return {"image": None, "video": None} |
| 131 | |
| 132 | def get_mm_max_tokens_per_item( |
| 133 | self, seq_len: int, mm_counts: Mapping[str, int] |
| 134 | ) -> Mapping[str, int] | None: |
| 135 | return super().get_mm_max_tokens_per_item(seq_len, mm_counts) |
| 136 | |
| 137 | |
| 138 | @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) |
| 139 | def _softcap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor: |
| 140 | # fp32 before tanh for numerical stability (matches HF DiffusionGemma). |
| 141 | # Compiling fuses the cast/div/tanh/mul into one elementwise kernel over |
| 142 | # the [num_tokens, vocab] logits instead of four separate passes. |
| 143 | logits = logits.float() |
| 144 | return torch.tanh(logits / cap) * cap |
| 145 | |
| 146 | |
| 147 | @MULTIMODAL_REGISTRY.register_processor( |
| 148 | Gemma4MultiModalProcessor, |
| 149 | info=DiffusionGemmaProcessingInfo, |
| 150 | dummy_inputs=Gemma4DummyInputsBuilder, |
| 151 | ) |
| 152 | class DiffusionGemmaForConditionalGeneration( |
| 153 | nn.Module, |
| 154 | SupportsMultiModal, |
| 155 | SupportsQuant, |
| 156 | ): |
| 157 | """DiffusionGemma for vLLM. |
| 158 | |
| 159 | Single Gemma4 backbone that switches between encoder and decoder mode. |
| 160 | The encoder path uses standard Gemma4 layers (causal attention, KV write). |
| 161 | The decoder path uses the same weights with bidirectional attention and |
| 162 | KV read-only, plus self-conditioning. |
| 163 | |
| 164 | Always includes a vision tower (same as Gemma4) for image understanding. |
| 165 | |
| 166 | In practice, the model's forward() dispatches based on the `mode` kwarg |
| 167 | set by DiffusionGemmaModelState.prepare_inputs(). |
| 168 | """ |
| 169 | |
| 170 | hf_to_vllm_mapper = WeightsMapper( |
| 171 | orig_to_new_prefix={ |
| 172 | "model.decoder.self_conditioning.": "self_conditioning.", |
| 173 | "model.decoder.": "model.", |
| 174 | "model.encoder.language_model.": "model.", |
| 175 | "model.encoder.vision_tower.": "vision_tower.", |
| 176 | "model.encoder.embed_vision.": "embed_vision.", |
| 177 | }, |
| 178 | ) |
| 179 | |
| 180 | packed_modules_mapping = { |
| 181 | "qkv_proj": ["q_proj", "k_proj", "v_proj"], |
| 182 | "gate_up_proj": ["gate_proj", "up_proj"], |
| 183 | } |
| 184 | |
| 185 | @staticmethod |
| 186 | def get_model_state_cls(): |
| 187 | return DiffusionGemmaModelState |
| 188 | |
| 189 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 190 | super().__init__() |
| 191 | config = vllm_config.model_config.hf_config |
| 192 | text_config = vllm_config.model_config.hf_text_config |
| 193 | self.config = config |
| 194 | self.model_dtype = vllm_config.model_config.dtype |
| 195 | |
| 196 | # DiffusionGemma's full-attention layers have NO v_proj — V is |
| 197 | # computed from k_proj's output (`value_states = key_states` before |
| 198 | # k_norm in `DiffusionGemmaDecoderTextAttention.forward`). This is |
| 199 | # the "k_eq_v" variant in our Gemma4 backbone. The checkpoint has no |
| 200 | # v_proj weights for full-attention layers; without this flag they |
| 201 | # would silently load with random V projections. |
| 202 | text_config.attention_k_eq_v = True |
| 203 | |
| 204 | # ---- Vision tower ---- |
| 205 | # Gemma4's image path, borrowed below, reads this flag. |
| 206 | lora_config = vllm_config.lora_config |
| 207 | self._enable_mm_lora = bool( |
| 208 | lora_config is not None and lora_config.enable_tower_connector_lora |
| 209 | ) |
| 210 | vision_config = getattr(config, "vision_config", None) |
| 211 | self.embed_vision: Gemma4MultimodalEmbedder | None |
| 212 | if vision_config is not None: |
| 213 | quant_config = vllm_config.quant_config |
| 214 | tower_quant: QuantizationConfig | None |
| 215 | if quant_config and quant_config.get_name() in [ |
| 216 | "bitsandbytes", |
| 217 | "torchao", |
| 218 | "compressed-tensors", |
| 219 | ]: |
| 220 | tower_quant = quant_config |
| 221 | else: |
| 222 | quantizable = ( |
| 223 | vision_config.hidden_size % 64 == 0 |
| 224 | and vision_config.intermediate_size % 64 == 0 |
| 225 | ) |
| 226 | tower_quant = quant_config if quantizable else None |
| 227 | |
| 228 | with self._mark_tower_model(vllm_config, {"image", "video"}): |
| 229 | self.vision_tower = AutoModel.from_config(config=vision_config) |
| 230 | self.embed_vision = Gemma4MultimodalEmbedder( |
| 231 | vision_config, |
| 232 | text_config, |
| 233 | quant_config=tower_quant, |
| 234 | prefix=maybe_prefix(prefix, "embed_vision"), |
| 235 | ) |
| 236 | recursive_replace_linear( |
| 237 | self.vision_tower, |
| 238 | tower_quant, |
| 239 | prefix=maybe_prefix(prefix, "vision_tower"), |
| 240 | ) |
| 241 | else: |
| 242 | self.vision_tower = None |
| 243 | self.embed_vision = None |
| 244 | |
| 245 | # ---- Language backbone (Gemma4Model) ---- |
| 246 | # Use maybe_prefix to ensure correct weight name prefixes for |
| 247 | # quantization. The quantization config uses hf_to_vllm_mapper to |
| 248 | # match checkpoint weight names to model parameter names. |
| 249 | self.model = Gemma4Model( |
| 250 | vllm_config=vllm_config, |
| 251 | prefix=maybe_prefix(prefix, "model"), |
| 252 | ) |
| 253 | |
| 254 | self.lm_head = ParallelLMHead( |
| 255 | num_embeddings=text_config.vocab_size, |
| 256 | embedding_dim=text_config.hidden_size, |
| 257 | quant_config=vllm_config.quant_config, |
| 258 | prefix=maybe_prefix(prefix, "lm_head"), |
| 259 | ) |
| 260 | |
| 261 | if text_config.tie_word_embeddings: |
| 262 | self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) |
| 263 | |
| 264 | # HF DiffusionGemma applies the final-logit softcap in fp32, before |
| 265 | # any other processing. Do it manually in `compute_logits` so the |
| 266 | # LogitsProcessor only handles the lm_head GEMM. |
| 267 | self.final_logit_softcapping = getattr( |
| 268 | text_config, "final_logit_softcapping", None |
| 269 | ) |
| 270 | self.logits_processor = LogitsProcessor( |
| 271 | text_config.vocab_size, |
| 272 | soft_cap=None, |
| 273 | ) |
| 274 | |
| 275 | sc_size = ( |
| 276 | getattr(config, "self_conditioning_size", None) |
| 277 | or text_config.intermediate_size |
| 278 | ) |
| 279 | self.self_conditioning = DiffusionGemmaSelfConditioning( |
| 280 | hidden_size=text_config.hidden_size, |
| 281 | self_conditioning_size=sc_size, |
| 282 | eps=getattr(text_config, "rms_norm_eps", 1e-6), |
| 283 | ) |
| 284 | |
| 285 | def compute_self_conditioning( |
| 286 | self, |
| 287 | inputs_embeds: torch.Tensor, |
| 288 | probs: torch.Tensor, |
| 289 | ) -> torch.Tensor: |
| 290 | embed_weight = self.model.embed_tokens.weight |
| 291 | soft_embeds = torch.matmul( |
| 292 | probs.to(embed_weight.dtype), embed_weight |
| 293 | ) * self.model.normalizer.to(inputs_embeds.dtype) |
| 294 | return self.self_conditioning(inputs_embeds, soft_embeds) |
| 295 | |
| 296 | # ------------------------------------------------------------------ # |
| 297 | # Multimodal: reuse Gemma4's image parsing, processing & embedding |
| 298 | # ------------------------------------------------------------------ # |
| 299 | # The vision tower, pooler, embed_vision, and their processing logic |
| 300 | # are architecturally identical to Gemma4. Delegate to avoid |
| 301 | # maintaining a duplicate copy. |
| 302 | |
| 303 | _parse_and_validate_image_input = ( |
| 304 | Gemma4ForConditionalGeneration._parse_and_validate_image_input |
| 305 | ) |
| 306 | _parse_and_validate_video_input = ( |
| 307 | Gemma4ForConditionalGeneration._parse_and_validate_video_input |
| 308 | ) |
| 309 | _parse_and_validate_multimodal_inputs = ( |
| 310 | Gemma4ForConditionalGeneration._parse_and_validate_multimodal_inputs |
| 311 | ) |
| 312 | _encoder_chunk = staticmethod(Gemma4ForConditionalGeneration._encoder_chunk) |
| 313 | _process_image_input = Gemma4ForConditionalGeneration._process_image_input |
| 314 | _process_video_input = Gemma4ForConditionalGeneration._process_video_input |
| 315 | embed_multimodal = Gemma4ForConditionalGeneration.embed_multimodal |
| 316 | |
| 317 | def get_mm_mapping(self) -> MultiModelKeys: |
| 318 | """Get the module prefix mapping for multimodal models.""" |
| 319 | return MultiModelKeys.from_string_field( |
| 320 | language_model="model", |
| 321 | connector=["embed_vision"], |
| 322 | tower_model=["vision_tower"], |
| 323 | ) |
| 324 | |
| 325 | # ------------------------------------------------------------------ # |
| 326 | # Forward |
| 327 | # ------------------------------------------------------------------ # |
| 328 | |
| 329 | def forward( |
| 330 | self, |
| 331 | input_ids: torch.Tensor, |
| 332 | positions: torch.Tensor, |
| 333 | intermediate_tensors: Any | None = None, |
| 334 | inputs_embeds: torch.Tensor | None = None, |
| 335 | **kwargs: Any, |
| 336 | ) -> torch.Tensor: |
| 337 | if intermediate_tensors is not None: |
| 338 | inputs_embeds = None |
| 339 | return self.model( |
| 340 | input_ids=input_ids, |
| 341 | positions=positions, |
| 342 | intermediate_tensors=intermediate_tensors, |
| 343 | inputs_embeds=inputs_embeds, |
| 344 | **kwargs, |
| 345 | ) |
| 346 | |
| 347 | def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: |
| 348 | states = getattr(self, "diffusion_states", None) |
| 349 | allowed = states.step_allowed if states is not None else None |
| 350 | if allowed is not None: |
| 351 | # One column per allowed id: a [rows, K] GEMM over the K gathered |
| 352 | # rows of the tied embedding. |
| 353 | logits = torch.nn.functional.linear( |
| 354 | hidden_states, self.lm_head.weight[allowed] |
| 355 | ).float() |
| 356 | if self.final_logit_softcapping is not None: |
| 357 | logits = _softcap_logits(logits, self.final_logit_softcapping) |
| 358 | return logits |
| 359 | logits = self.logits_processor(self.lm_head, hidden_states) |
| 360 | if logits is not None and self.final_logit_softcapping is not None: |
| 361 | logits = _softcap_logits(logits, self.final_logit_softcapping) |
| 362 | return logits |
| 363 | |
| 364 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 365 | # Some checkpoints carry a vestigial Gemma3n-style embedding table. |
| 366 | loader = AutoWeightsLoader( |
| 367 | self, ignore_unexpected_prefixes=["embed_vision.embedding."] |
| 368 | ) |
| 369 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 370 | |
| 371 | @classmethod |
| 372 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 373 | if modality == "image": |
| 374 | return "<image_soft_token>" |
| 375 | if modality == "video": |
| 376 | return "<|video|>" |
| 377 | raise ValueError(f"Unsupported modality: {modality}") |
| 378 | |
| 379 | |
| 380 | @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) |
| 381 | def _compute_num_rejected( |
| 382 | num_logits: torch.Tensor, |
| 383 | num_sampled: torch.Tensor, |
| 384 | query_start_loc: torch.Tensor, |
| 385 | ) -> torch.Tensor: |
| 386 | query_lens = query_start_loc[1:] - query_start_loc[:-1] |
| 387 | num_rejected = num_logits - num_sampled |
| 388 | is_denoise = (num_logits > 0) & (num_sampled == 0) |
| 389 | return torch.where(is_denoise, query_lens, num_rejected) |
| 390 | |
| 391 | |
| 392 | def _concat_logprob_stashes( |
| 393 | parts: list[LogprobsTensors], cu_num_generated_tokens: list[int] |
| 394 | ) -> LogprobsTensors: |
| 395 | """Join the logprobs stashed for the requests committing this step. |
| 396 | |
| 397 | Each stash is as wide as the widest logprobs request in the batch at the |
| 398 | step that request converged, so stashes from different steps can differ |
| 399 | in width. Pad the narrow ones the way compute_topk_scores pads a mixed |
| 400 | batch: token id 0 at -inf, which the output processor never reports. |
| 401 | """ |
| 402 | width = max(p.logprob_token_ids.shape[1] for p in parts) |
| 403 | |
| 404 | def pad(t: torch.Tensor, value: float) -> torch.Tensor: |
| 405 | return F.pad(t, (0, width - t.shape[1]), value=value) |
| 406 | |
| 407 | return LogprobsTensors( |
| 408 | logprob_token_ids=torch.cat([pad(p.logprob_token_ids, 0) for p in parts]), |
| 409 | logprobs=torch.cat([pad(p.logprobs, float("-inf")) for p in parts]), |
| 410 | selected_token_ranks=torch.cat([p.selected_token_ranks for p in parts]), |
| 411 | cu_num_generated_tokens=cu_num_generated_tokens, |
| 412 | ) |
| 413 | |
| 414 | |
| 415 | _MASKED_LOGIT = -1e20 |
| 416 | |
| 417 | |
| 418 | def _mask_rows_to_allowed( |
| 419 | logits: torch.Tensor, |
| 420 | row_starts: Sequence[int] | np.ndarray, |
| 421 | row_lens: Sequence[int] | np.ndarray, |
| 422 | allowed_per_row: Sequence[torch.Tensor | None], |
| 423 | ) -> torch.Tensor: |
| 424 | """Mask every column outside a request's allowed ids on that request's |
| 425 | rows. Request i owns rows [row_starts[i], +row_lens[i]); None leaves its |
| 426 | rows alone. Returns a copy when any row is masked, so the caller's tensor |
| 427 | (possibly the runner's) is never written. |
| 428 | |
| 429 | Softmax over a masked row equals the K-space distribution the shared |
| 430 | fast path computes, so both paths give the same reads. The mask value is |
| 431 | a large finite negative rather than -inf: the entropy is probs times |
| 432 | log-probs, and 0 * -inf is NaN, while 0 * -1e20 is 0. It stays finite |
| 433 | after the greedy temperature clamp (1e-10) scales it by 1e10. |
| 434 | """ |
| 435 | out: torch.Tensor | None = None |
| 436 | keep = torch.zeros(logits.shape[-1], dtype=torch.bool, device=logits.device) |
| 437 | for start, n, allowed in zip(row_starts, row_lens, allowed_per_row): |
| 438 | if allowed is None: |
| 439 | continue |
| 440 | if out is None: |
| 441 | out = logits.clone() |
| 442 | keep.zero_() |
| 443 | keep[allowed] = True |
| 444 | out[int(start) : int(start) + int(n)].masked_fill_(~keep, _MASKED_LOGIT) |
| 445 | return logits if out is None else out |
| 446 | |
| 447 | |
| 448 | # Tiles specialize the graph on batch size 1, on a tile as wide as the canvas, |
| 449 | # on compute_sc and on sizes that coincide with the state buffers. That set |
| 450 | # is small but passes Dynamo's default of 8, after which every step would run |
| 451 | # eager: about twice as slow for a self-conditioned step. |
| 452 | @torch._dynamo.config.patch(recompile_limit=64) |
| 453 | @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) |
| 454 | def _denoise_temperature( |
| 455 | step_tensor: torch.Tensor, |
| 456 | slots: torch.Tensor, |
| 457 | max_denoising_steps: float, |
| 458 | t_min: float, |
| 459 | t_max: float, |
| 460 | ) -> torch.Tensor: |
| 461 | """The schedule's temperature for each slot at its current step.""" |
| 462 | steps_f = step_tensor[slots].float() |
| 463 | remaining = (max_denoising_steps - steps_f).clamp(min=1.0) |
| 464 | return t_min + (t_max - t_min) * (remaining / max_denoising_steps) |
| 465 | |
| 466 | |
| 467 | def _compiled_sample_step( |
| 468 | # Per-position statistics of the temperature-scaled logits, from |
| 469 | # sample_row_stats: [num_decode, CL] each, and the softmax |
| 470 | # [num_decode, CL, vocab] in the embedding dtype when compute_sc. |
| 471 | new_tokens: torch.Tensor, |
| 472 | argmax_tokens: torch.Tensor, |
| 473 | token_entropy: torch.Tensor, |
| 474 | probs: torch.Tensor | None, |
| 475 | # Request mapping |
| 476 | decode_slots: torch.Tensor, # [num_decode] int64 → slot indices |
| 477 | decode_idx: torch.Tensor, # [num_decode] int64 → position in num_reqs |
| 478 | all_slots: torch.Tensor, # [num_reqs] int64 → all slot indices |
| 479 | valid_canvas_len: torch.Tensor, # [num_decode] int64 → real canvas length (<=CL) |
| 480 | # State tensors (modified in-place) |
| 481 | canvas: torch.Tensor, # [max_num_reqs, CL] |
| 482 | argmax_canvas: torch.Tensor, # [max_num_reqs, CL] |
| 483 | step_tensor: torch.Tensor, # [max_num_reqs] |
| 484 | is_encoder_phase: torch.Tensor, # [max_num_reqs] |
| 485 | confident_tensor: torch.Tensor, # [max_num_reqs] |
| 486 | sc_embeds: torch.Tensor, # [max_num_reqs, CL, hidden] |
| 487 | embed_weight: torch.Tensor, # [vocab, hidden] |
| 488 | normalizer: torch.Tensor, |
| 489 | history: torch.Tensor, # [max_num_reqs, ST, CL] |
| 490 | history_len_tensor: torch.Tensor, # [max_num_reqs] |
| 491 | max_steps_tensor: torch.Tensor, # [max_num_reqs] int32, per-slot step cap |
| 492 | pin_mask: torch.Tensor, # [max_num_reqs, CL] bool, positions held at the seed |
| 493 | seed_canvas: torch.Tensor, # [max_num_reqs, CL] |
| 494 | read_only: torch.Tensor, # [max_num_reqs] bool, emit without a commit forward |
| 495 | # Output tensors (modified in-place) |
| 496 | sampled: torch.Tensor, # [num_reqs, CL] |
| 497 | num_sampled: torch.Tensor, # [num_reqs] |
| 498 | draft_tokens: torch.Tensor, # [max_num_reqs, >=CL] |
| 499 | # Scalar config |
| 500 | confidence_threshold: float, |
| 501 | vocab_size: int, |
| 502 | CL: int, |
| 503 | ST: int, |
| 504 | # Sampler config |
| 505 | entropy_bound: float, |
| 506 | # Tensor-parallel vocab sharding for the self-conditioning matmul. |
| 507 | # ``embed_weight`` is vocab-sharded ([vocab/tp, hidden]) while ``probs`` |
| 508 | # spans the full vocab; [sc_vocab_start, sc_vocab_end) is this rank's slice. |
| 509 | sc_vocab_start: int, |
| 510 | sc_vocab_end: int, |
| 511 | tp_size: int, |
| 512 | tp_group_name: str, |
| 513 | compute_sc: bool = True, |
| 514 | ) -> None: |
| 515 | """Compiled decode step: confidence → accept/renoise → convergence, as |
| 516 | vectorized PyTorch ops over [num_decode, CL] tensors. The per-position |
| 517 | statistics (argmax, Gumbel-max sample, entropy, softmax) come from one |
| 518 | pass over the logits in ``sample_row_stats``.""" |
| 519 | num_decode = decode_slots.shape[0] |
| 520 | device = decode_slots.device |
| 521 | |
| 522 | # ---- Phase 3: Confidence ---- |
| 523 | # A canvas truncated near max_model_len is zero-padded up to CL by the |
| 524 | # caller; those padded rows are uniform (max entropy, argmax 0), so they |
| 525 | # never trigger early convergence and are stable, and only the real |
| 526 | # ``valid_canvas_len`` tokens are committed (num_sampled below). |
| 527 | mean_entropy = token_entropy.mean(dim=-1) # [num_decode] |
| 528 | confident_tensor[decode_slots] = mean_entropy < confidence_threshold |
| 529 | |
| 530 | # ---- Phase 4: Entropy-bound acceptance mask ---- |
| 531 | sorted_ent, sorted_idx = torch.sort(token_entropy, dim=-1) |
| 532 | cumsum_ent = torch.cumsum(sorted_ent, dim=-1) |
| 533 | cummax_ent = torch.cummax(sorted_ent, dim=-1).values |
| 534 | sorted_mask = (cumsum_ent - cummax_ent) <= entropy_bound |
| 535 | eb_mask = torch.zeros_like(sorted_mask) |
| 536 | eb_mask.scatter_(1, sorted_idx, sorted_mask) |
| 537 | |
| 538 | # ---- Phase 5: Post-sample ---- |
| 539 | is_commit = is_encoder_phase[decode_slots] # [num_decode] |
| 540 | is_denoise = ~is_commit |
| 541 | cur_step = step_tensor[decode_slots].float() |
| 542 | |
| 543 | # Step update: +1 for denoise, reset to 0 for commit |
| 544 | new_step_val = torch.where( |
| 545 | is_denoise, |
| 546 | (cur_step + 1).to(step_tensor.dtype), |
| 547 | step_tensor.new_zeros(num_decode), |
| 548 | ) |
| 549 | step_tensor[decode_slots] = new_step_val |
| 550 | |
| 551 | # Random tokens for renoise / canvas reinit |
| 552 | # Renoise draws from the whole vocabulary even for a constrained read. |
| 553 | # The model was trained on random-vocabulary noise. Noise drawn from the |
| 554 | # allowed set looks like garbled text, and later steps denoise it into |
| 555 | # more garbled text. |
| 556 | random_tokens = torch.randint( |
| 557 | 0, vocab_size, (num_decode, CL), device=device, dtype=canvas.dtype |
| 558 | ) |
| 559 | |
| 560 | # Compute denoise canvas (accept/renoise). Pinned positions keep the seed. |
| 561 | denoise_canvas = torch.where(eb_mask, new_tokens, random_tokens) |
| 562 | denoise_canvas = torch.where( |
| 563 | pin_mask[decode_slots], seed_canvas[decode_slots], denoise_canvas |
| 564 | ) |
| 565 | |
| 566 | # Canvas: commit → random reinit, denoise → accept/renoise result |
| 567 | canvas[decode_slots] = torch.where( |
| 568 | is_commit.unsqueeze(1), random_tokens, denoise_canvas |
| 569 | ) |
| 570 | |
| 571 | # History: write argmax_tokens for denoise requests at circular position |
| 572 | hist_len = history_len_tensor[decode_slots] |
| 573 | write_pos = hist_len % ST |
| 574 | for i in range(ST): |
| 575 | write_here = ((write_pos == i) & is_denoise).unsqueeze(1) |
| 576 | history[decode_slots, i] = torch.where( |
| 577 | write_here, argmax_tokens, history[decode_slots, i] |
| 578 | ) |
| 579 | |
| 580 | # Argmax canvas: update for denoise, preserve for commit |
| 581 | argmax_canvas[decode_slots] = torch.where( |
| 582 | is_denoise.unsqueeze(1), argmax_tokens, argmax_canvas[decode_slots] |
| 583 | ) |
| 584 | |
| 585 | # History length: increment for denoise, reset for commit |
| 586 | new_hist_len = torch.where(is_denoise, hist_len + 1, hist_len.new_zeros(num_decode)) |
| 587 | history_len_tensor[decode_slots] = new_hist_len |
| 588 | |
| 589 | # ---- Phase 6: Stability + convergence ---- |
| 590 | ref = history[decode_slots, 0] |
| 591 | mismatch = torch.zeros(num_decode, device=device, dtype=torch.int32) |
| 592 | for h in range(1, ST): |
| 593 | mismatch = mismatch + (ref != history[decode_slots, h]).sum(dim=-1).int() |
| 594 | stable = mismatch == 0 |
| 595 | |
| 596 | step_after = step_tensor[decode_slots] |
| 597 | converged = (stable & confident_tensor[decode_slots] & (new_hist_len >= ST)) | ( |
| 598 | step_after >= max_steps_tensor[decode_slots] |
| 599 | ) |
| 600 | # Commit done → denoise next (False); denoise converged → commit next (True) |
| 601 | is_encoder_phase[decode_slots] = torch.where( |
| 602 | is_commit, is_commit.new_zeros(num_decode), converged |
| 603 | ) |
| 604 | |
| 605 | emit = is_commit | (is_denoise & converged & read_only[decode_slots]) |
| 606 | sampled[decode_idx] = argmax_canvas[decode_slots].to(sampled.dtype) * emit[:, None] |
| 607 | num_sampled[decode_idx] = (emit * valid_canvas_len).to(num_sampled.dtype) |
| 608 | |
| 609 | # SC soft embedding: store ``probs @ embed_weight`` (the value the next step's |
| 610 | # self-conditioning MLP consumes) only for slots that will denoise next — i.e. |
| 611 | # this step denoised AND it isn't about to commit (is_encoder_phase now False). |
| 612 | # Masking here (rather than in the consumer) lets _apply_self_conditioning read |
| 613 | # sc_embeds directly. Storing the [.., hidden] soft embed instead of the full |
| 614 | # [.., vocab] probs avoids a giant persistent buffer. |
| 615 | sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None] |
| 616 | if compute_sc: |
| 617 | assert probs is not None |
| 618 | # Self-conditioning soft embed = probs @ embed_tokens.weight. Under |
| 619 | # tensor parallelism the embedding is vocab-sharded ([vocab/tp, |
| 620 | # hidden]) while probs spans the full vocab, so each rank multiplies |
| 621 | # its local vocab slice [sc_vocab_start, sc_vocab_end) and the |
| 622 | # partials are summed across ranks. |
| 623 | local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype) |
| 624 | soft_embeds = torch.matmul( |
| 625 | local_probs, embed_weight[: sc_vocab_end - sc_vocab_start] |
| 626 | ) |
| 627 | if tp_size > 1: |
| 628 | soft_embeds = torch.ops.vllm.all_reduce( |
| 629 | soft_embeds, group_name=tp_group_name |
| 630 | ) |
| 631 | soft_embeds = soft_embeds * normalizer |
| 632 | # A pinned position holds its seed token as input. Zero its soft embed |
| 633 | # too, or the model's own prediction there reaches the next step |
| 634 | # through self-conditioning. The buffer is fp32 while the embedding is |
| 635 | # in the model dtype: compiled code casts on the store but eager does |
| 636 | # not, and the step runs eager once torch.compile hits its recompile |
| 637 | # limit. |
| 638 | sc_pin = (~pin_mask[decode_slots]).unsqueeze(-1) |
| 639 | sc_embeds[decode_slots] = (soft_embeds * sc_keep * sc_pin).to(sc_embeds.dtype) |
| 640 | else: |
| 641 | # Every slot in this tile ends after this step, so the soft embed |
| 642 | # would never be read. The matmul is a full pass over the vocabulary |
| 643 | # matrix, so skip it. |
| 644 | sc_embeds[decode_slots] = 0 |
| 645 | |
| 646 | # Overwrite canvas with argmax for newly converged denoise requests |
| 647 | newly_converged = (converged & is_denoise).unsqueeze(1) |
| 648 | canvas[decode_slots] = torch.where( |
| 649 | newly_converged, argmax_canvas[decode_slots], canvas[decode_slots] |
| 650 | ) |
| 651 | is_encoder_phase[decode_slots] &= ~read_only[decode_slots] |
| 652 | |
| 653 | # ---- Phase 7: Copy canvas → draft_tokens for all slots ---- |
| 654 | draft_tokens[all_slots, :CL] = canvas[all_slots] |
| 655 | |
| 656 | |
| 657 | class DiffusionGemmaRequestStates: |
| 658 | """Pre-allocated GPU tensors for DiffusionGemma per-request state. |
| 659 | |
| 660 | Follows the indexed-slot pattern used by ``RequestState``. |
| 661 | """ |
| 662 | |
| 663 | def __init__( |
| 664 | self, |
| 665 | max_num_reqs: int, |
| 666 | canvas_length: int, |
| 667 | vocab_size: int, |
| 668 | max_denoising_steps: int, |
| 669 | device: torch.device, |
| 670 | hidden_size: int, |
| 671 | stability_threshold: int, |
| 672 | ): |
| 673 | self.max_num_reqs = max_num_reqs |
| 674 | self.canvas_length = canvas_length |
| 675 | self.vocab_size = vocab_size |
| 676 | self.max_denoising_steps = max_denoising_steps |
| 677 | self.stability_threshold = stability_threshold |
| 678 | self.device = device |
| 679 | |
| 680 | self.is_encoder_phase = torch.zeros( |
| 681 | max_num_reqs, dtype=torch.bool, device=device |
| 682 | ) |
| 683 | # Canvas tokens [max_num_reqs, canvas_length] |
| 684 | self.canvas = torch.zeros( |
| 685 | max_num_reqs, canvas_length, dtype=torch.int64, device=device |
| 686 | ) |
| 687 | # Step counter (counts up from 0 to max_denoising_steps) |
| 688 | self.step = torch.zeros( |
| 689 | max_num_reqs, |
| 690 | dtype=torch.int32, |
| 691 | device=device, |
| 692 | ) |
| 693 | # Accepted canvas history for stability check |
| 694 | self.accepted_canvas_history = torch.zeros( |
| 695 | max_num_reqs, |
| 696 | stability_threshold, |
| 697 | canvas_length, |
| 698 | dtype=torch.int64, |
| 699 | device=device, |
| 700 | ) |
| 701 | self.accepted_canvas_history_len = torch.zeros( |
| 702 | max_num_reqs, dtype=torch.int32, device=device |
| 703 | ) |
| 704 | # Latest argmax(processed_logits) per slot — what we COMMIT. |
| 705 | # NOT `current_canvas` (which is the post-renoise stochastic input for |
| 706 | # the next denoise step). We keep this separate from `canvas` because |
| 707 | # canvas gets renoised in-place during denoise, while argmax_canvas is |
| 708 | # the deterministic best-guess we ultimately emit. |
| 709 | self.argmax_canvas = torch.zeros( |
| 710 | max_num_reqs, canvas_length, dtype=torch.int64, device=device |
| 711 | ) |
| 712 | |
| 713 | # Per-slot prompt length (set by add_request). |
| 714 | self.prompt_len = torch.zeros( |
| 715 | max_num_reqs, |
| 716 | dtype=torch.int32, |
| 717 | device=device, |
| 718 | ) |
| 719 | |
| 720 | # Per-slot confidence flag, set by the sampler each step. |
| 721 | self.confident = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) |
| 722 | |
| 723 | # Per-slot step cap, lowered per request from extra_args. |
| 724 | self.max_steps = torch.full( |
| 725 | (max_num_reqs,), max_denoising_steps, dtype=torch.int32, device=device |
| 726 | ) |
| 727 | # Seed canvases replace the random canvas after prefill. The host-side |
| 728 | # sets gate the seed and read-only work so plain generation skips it. |
| 729 | self.seed_canvas = torch.zeros( |
| 730 | max_num_reqs, canvas_length, dtype=torch.int64, device=device |
| 731 | ) |
| 732 | self.has_seed = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) |
| 733 | self.seeded_slots: set[int] = set() |
| 734 | # Pinned positions keep their seed value on every denoise step, so a |
| 735 | # multi-step read denoises the canvas the request seeded. |
| 736 | self.pin_mask = torch.zeros( |
| 737 | max_num_reqs, canvas_length, dtype=torch.bool, device=device |
| 738 | ) |
| 739 | # Read-only slots emit on their converging step and skip the commit |
| 740 | # forward. |
| 741 | self.read_only = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) |
| 742 | # Constrained reads: slot -> the allowed ids (the request's |
| 743 | # logprob_token_ids). A step whose decode slots all share one set |
| 744 | # runs the unembedding, sampler and self-conditioning over that set. |
| 745 | self.constrained: dict[int, tuple[int, ...]] = {} |
| 746 | self.step_allowed: torch.Tensor | None = None |
| 747 | self._allowed_cache: dict[tuple[int, ...], torch.Tensor] = {} |
| 748 | self.read_only_slots: set[int] = set() |
| 749 | # Slots capped at one denoise step never consume a soft embed. |
| 750 | self.single_step_slots: set[int] = set() |
| 751 | # Per-slot canvas width, at most canvas_length. The scheduler schedules |
| 752 | # this many draft tokens for the slot and the sampler pads the rest. |
| 753 | self.canvas_width_np = np.full(max_num_reqs, canvas_length, dtype=np.int32) |
| 754 | |
| 755 | # Per-slot self-conditioning soft embedding (probs @ embed_weight) from |
| 756 | # the previous denoise step. Storing the [.., hidden] soft embed instead |
| 757 | # of the full [.., vocab] distribution shrinks this buffer by |
| 758 | # vocab/hidden (~170x) and moves the matmul to denoise time; the result |
| 759 | # is identical (SC consumes probs @ embed_weight anyway). |
| 760 | self.self_conditioning_embeds = torch.zeros( |
| 761 | max_num_reqs, canvas_length, hidden_size, dtype=torch.float32, device=device |
| 762 | ) |
| 763 | |
| 764 | def init_canvas(self, slot_indices: torch.Tensor) -> None: |
| 765 | """Initialize canvas with random tokens for the given slots. |
| 766 | |
| 767 | `slot_indices` must already be on device to avoid a cpu->gpu sync. |
| 768 | """ |
| 769 | n = slot_indices.shape[0] |
| 770 | self.canvas[slot_indices] = torch.randint( |
| 771 | 0, |
| 772 | self.vocab_size, |
| 773 | (n, self.canvas_length), |
| 774 | dtype=torch.int64, |
| 775 | device=self.device, |
| 776 | ) |
| 777 | |
| 778 | def add_request(self, slot_idx: int) -> None: |
| 779 | self.is_encoder_phase[slot_idx].fill_(True) |
| 780 | self.init_canvas(async_tensor_h2d([slot_idx], device=self.device)) |
| 781 | self.step[slot_idx].fill_(0) |
| 782 | self.accepted_canvas_history_len[slot_idx].fill_(0) |
| 783 | self.self_conditioning_embeds[slot_idx] = 0 |
| 784 | self.max_steps[slot_idx].fill_(self.max_denoising_steps) |
| 785 | self.has_seed[slot_idx].fill_(False) |
| 786 | self.seeded_slots.discard(slot_idx) |
| 787 | self.pin_mask[slot_idx].fill_(False) |
| 788 | self.read_only[slot_idx].fill_(False) |
| 789 | self.read_only_slots.discard(slot_idx) |
| 790 | self.single_step_slots.discard(slot_idx) |
| 791 | self.constrained.pop(slot_idx, None) |
| 792 | self.canvas_width_np[slot_idx] = self.canvas_length |
| 793 | |
| 794 | def remove_request(self, slot_idx: int) -> None: |
| 795 | # add_request resets the GPU flags before a slot is reused. The host |
| 796 | # sets gate work on every step, so they clear here. |
| 797 | self.is_encoder_phase[slot_idx].fill_(False) |
| 798 | self.accepted_canvas_history_len[slot_idx].fill_(0) |
| 799 | self.self_conditioning_embeds[slot_idx] = 0 |
| 800 | self.seeded_slots.discard(slot_idx) |
| 801 | self.read_only_slots.discard(slot_idx) |
| 802 | self.single_step_slots.discard(slot_idx) |
| 803 | self.constrained.pop(slot_idx, None) |
| 804 | |
| 805 | def set_seed_canvas(self, slot_idx: int, ids: list[int]) -> None: |
| 806 | """``ids`` covers the slot's canvas width; positions past it are never |
| 807 | scheduled.""" |
| 808 | self.seed_canvas[slot_idx, : len(ids)] = async_tensor_h2d( |
| 809 | ids, dtype=torch.int64, device=self.device |
| 810 | ) |
| 811 | self.has_seed[slot_idx].fill_(True) |
| 812 | self.seeded_slots.add(slot_idx) |
| 813 | |
| 814 | def set_pins(self, slot_idx: int, positions: list[int]) -> None: |
| 815 | """Hold ``positions`` of the slot's seed canvas through every denoise |
| 816 | step. Validated against the request's canvas width upstream.""" |
| 817 | self.pin_mask[slot_idx].fill_(False) |
| 818 | self.pin_mask[ |
| 819 | slot_idx, async_tensor_h2d(positions, dtype=torch.int64, device=self.device) |
| 820 | ] = True |
| 821 | |
| 822 | def allowed_tensor(self, ids: tuple[int, ...]) -> torch.Tensor: |
| 823 | """``ids`` as an int64 tensor on the device, built once per tuple.""" |
| 824 | t = self._allowed_cache.get(ids) |
| 825 | if t is None: |
| 826 | t = torch.tensor(ids, dtype=torch.int64, device=self.device) |
| 827 | if len(self._allowed_cache) > 64: |
| 828 | self._allowed_cache.clear() |
| 829 | self._allowed_cache[ids] = t |
| 830 | return t |
| 831 | |
| 832 | def batch_allowed(self, slots: list[int]) -> torch.Tensor | None: |
| 833 | """The allowed ids shared by every one of ``slots``, or None. |
| 834 | |
| 835 | None means the step runs over the full vocabulary. The sampler then |
| 836 | masks each constrained slot's logit rows to its own set, so a |
| 837 | constrained request reads the same way whoever shares its batch. |
| 838 | """ |
| 839 | if not slots or not self.constrained: |
| 840 | return None |
| 841 | first = self.constrained.get(slots[0]) |
| 842 | if first is None or any(self.constrained.get(s) != first for s in slots[1:]): |
| 843 | return None |
| 844 | return self.allowed_tensor(first) |
| 845 | |
| 846 | def set_read_only(self, slot_idx: int) -> None: |
| 847 | self.read_only[slot_idx].fill_(True) |
| 848 | self.read_only_slots.add(slot_idx) |
| 849 | |
| 850 | def apply_seed_canvases( |
| 851 | self, slots_np: np.ndarray, slots_gpu: torch.Tensor |
| 852 | ) -> None: |
| 853 | """Replace the canvas of every seeded slot among ``slots_gpu``.""" |
| 854 | if self.seeded_slots.isdisjoint(slots_np.tolist()): |
| 855 | return |
| 856 | self.canvas[slots_gpu] = torch.where( |
| 857 | self.has_seed[slots_gpu, None], |
| 858 | self.seed_canvas[slots_gpu], |
| 859 | self.canvas[slots_gpu], |
| 860 | ) |
| 861 | |
| 862 | |
| 863 | class DiffusionGemmaModelState(ModelState): |
| 864 | """ModelState for DiffusionGemma. |
| 865 | |
| 866 | Single Gemma4 backbone in two modes: |
| 867 | - encoder mode (num_draft_tokens == 0): causal attention, writes KV |
| 868 | - decoder mode (num_draft_tokens > 0): bidirectional attention, reads KV |
| 869 | """ |
| 870 | |
| 871 | def __init__( |
| 872 | self, |
| 873 | vllm_config: VllmConfig, |
| 874 | model: nn.Module, |
| 875 | encoder_cache: Any, |
| 876 | device: torch.device, |
| 877 | ) -> None: |
| 878 | super().__init__(vllm_config, model, encoder_cache, device) |
| 879 | |
| 880 | # Per-step MM data produced by prepare_inputs_embeds and consumed by |
| 881 | # prepare_inputs. Stored as raw (mm_embeds, is_mm_embed) so that |
| 882 | # prepare_inputs can call embed_input_ids directly into the |
| 883 | # persistent _inputs_embeds_buf, avoiding the intermediate copy |
| 884 | # through encoder_runner.inputs_embeds. |
| 885 | self._pending_mm_embeds: tuple[list[torch.Tensor], torch.Tensor] | None = None |
| 886 | |
| 887 | diffusion_config = vllm_config.diffusion_config |
| 888 | canvas_length = diffusion_config.canvas_length if diffusion_config else 32 |
| 889 | |
| 890 | text_config = self.model_config.hf_text_config |
| 891 | self.gen_config = self.model_config.try_get_generation_config() |
| 892 | max_denoising_steps = ( |
| 893 | diffusion_config.max_denoising_steps if diffusion_config else None |
| 894 | ) or self.gen_config.get("max_denoising_steps", 48) |
| 895 | self.diffusion_states = DiffusionGemmaRequestStates( |
| 896 | max_num_reqs=self.max_num_reqs, |
| 897 | canvas_length=canvas_length, |
| 898 | vocab_size=self.model_config.get_vocab_size(), |
| 899 | max_denoising_steps=max_denoising_steps, |
| 900 | device=device, |
| 901 | hidden_size=text_config.hidden_size, |
| 902 | # In Transformers, `stability_threshold=1` (the default) means the current |
| 903 | # step must match the previous step. In vLLM, the history buffer includes |
| 904 | # the current step, so we add 1 to match the same behavior. |
| 905 | stability_threshold=self.gen_config["stability_threshold"] + 1, |
| 906 | ) |
| 907 | # compute_logits reads the step's shared allowed set from here. |
| 908 | self.model.diffusion_states = self.diffusion_states |
| 909 | self._req_id_to_index: dict[str, int] = {} |
| 910 | |
| 911 | # Persistent buffer for per-request causal flags, updated in-place |
| 912 | # so FULL CUDA graph replay sees the latest values. Must be int32: |
| 913 | # FlashAttentionMetadataBuilder.build() rejects other dtypes, since |
| 914 | # an out-of-place cast there would detach the captured graph from |
| 915 | # this buffer and freeze replay at the capture-time snapshot. |
| 916 | self._causal_buf = torch.zeros( |
| 917 | self.max_num_reqs, dtype=torch.int32, device=device |
| 918 | ) |
| 919 | |
| 920 | # Persistent inputs_embeds buffer — required so FULL CUDA graph |
| 921 | # capture and runtime point at the SAME memory address. |
| 922 | # `prepare_dummy_inputs` (capture path) and `prepare_inputs` (runtime |
| 923 | # path) both must hand the captured graph a tensor at this address. |
| 924 | self._inputs_embeds_buf = torch.zeros( |
| 925 | self.max_num_tokens, |
| 926 | text_config.hidden_size, |
| 927 | dtype=self.model_config.dtype, |
| 928 | device=device, |
| 929 | ) |
| 930 | |
| 931 | def get_supported_generation_tasks(self): |
| 932 | return ("generate",) |
| 933 | |
| 934 | def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: |
| 935 | diffusion_config = self.vllm_config.diffusion_config |
| 936 | gen = self.gen_config |
| 937 | sampler_cfg = gen.get("sampler_config") or {} |
| 938 | if "EntropyBound" not in sampler_cfg.get("_cls_name", ""): |
| 939 | raise ValueError("DiffusionGemma requires an EntropyBound sampler_config") |
| 940 | entropy_bound = sampler_cfg.get("entropy_bound") |
| 941 | if entropy_bound is None or entropy_bound <= 0: |
| 942 | raise ValueError( |
| 943 | f"entropy_bound must be a positive float (got {entropy_bound})" |
| 944 | ) |
| 945 | # The self-conditioning matmul (probs @ embed_tokens.weight) runs over a |
| 946 | # vocab-parallel embedding shard. Hand the sampler this rank's vocab |
| 947 | # slice and TP group so it can all-reduce the partial products. |
| 948 | embed_tokens = self.model.model.embed_tokens |
| 949 | shard = embed_tokens.shard_indices |
| 950 | tp_group = get_tp_group() |
| 951 | return DiffusionSampler( |
| 952 | sampler=sampler, |
| 953 | diffusion_config=diffusion_config, |
| 954 | vocab_size=self.model_config.get_vocab_size(), |
| 955 | diffusion_states=self.diffusion_states, |
| 956 | t_min=gen["t_min"], |
| 957 | t_max=gen["t_max"], |
| 958 | entropy_bound=entropy_bound, |
| 959 | confidence_threshold=gen["confidence_threshold"], |
| 960 | embed_weight=embed_tokens.weight, |
| 961 | normalizer=self.model.model.normalizer, |
| 962 | sc_vocab_start=shard.org_vocab_start_index, |
| 963 | sc_vocab_end=shard.org_vocab_end_index, |
| 964 | tp_size=tp_group.world_size, |
| 965 | tp_group_name=tp_group.unique_name, |
| 966 | ), None |
| 967 | |
| 968 | def apply_staged_writes(self) -> None: |
| 969 | pass |
| 970 | |
| 971 | def add_request(self, req_index: int, new_req_data: Any) -> None: |
| 972 | self._req_id_to_index[new_req_data.req_id] = req_index |
| 973 | self.diffusion_states.add_request(req_index) |
| 974 | if not new_req_data.req_id.startswith("_warmup_"): |
| 975 | prompt_len = len(new_req_data.prompt_token_ids) |
| 976 | self.diffusion_states.prompt_len[req_index].fill_(prompt_len) |
| 977 | |
| 978 | def remove_request(self, req_id: str) -> None: |
| 979 | idx = self._req_id_to_index.pop(req_id, None) |
| 980 | if idx is not None: |
| 981 | self.diffusion_states.remove_request(idx) |
| 982 | |
| 983 | def prepare_inputs_embeds( |
| 984 | self, |
| 985 | scheduled_encoder_inputs: dict[str, list[int]], |
| 986 | input_batch: InputBatch, |
| 987 | req_states: RequestState, |
| 988 | ) -> torch.Tensor | None: |
| 989 | if not self.supports_mm_inputs: |
| 990 | return None |
| 991 | |
| 992 | mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs( |
| 993 | scheduled_encoder_inputs |
| 994 | ) |
| 995 | if mm_kwargs: |
| 996 | encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs) |
| 997 | self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs)) |
| 998 | |
| 999 | mm_embeds, is_mm_embed = self.gather_mm_embeddings(input_batch) |
| 1000 | |
| 1001 | if not mm_embeds: |
| 1002 | # No MM tokens in this batch (e.g. all-decode step). |
| 1003 | # prepare_inputs will use embed_input_ids (text-only) directly. |
| 1004 | self._pending_mm_embeds = None |
| 1005 | return None |
| 1006 | |
| 1007 | # Stash raw MM ingredients for prepare_inputs to merge directly |
| 1008 | # into the persistent buffer, avoiding the intermediate copy |
| 1009 | # through encoder_runner.inputs_embeds. |
| 1010 | self._pending_mm_embeds = (mm_embeds, is_mm_embed) |
| 1011 | return None |
| 1012 | |
| 1013 | def _apply_self_conditioning( |
| 1014 | self, |
| 1015 | decode_slots_np: np.ndarray, |
| 1016 | decode_idx_np: np.ndarray, |
| 1017 | query_start_loc_np: np.ndarray, |
| 1018 | inputs_embeds: torch.Tensor, |
| 1019 | sc_embeds: torch.Tensor, |
| 1020 | ) -> None: |
| 1021 | # One self-conditioning MLP call per decode request, over that request's |
| 1022 | # query span [start, end) = its canvas. The span is the full canvas (CL) |
| 1023 | # or, for the final canvas truncated near max_model_len, fewer than CL |
| 1024 | # positions. sc_embeds already holds probs @ embed_weight from the prior |
| 1025 | # denoise step, masked to zero by the sampler for slots not denoising |
| 1026 | # this step; only the MLP runs here. CPU metadata -> no GPU syncs. |
| 1027 | for slot, idx in zip(decode_slots_np.tolist(), decode_idx_np.tolist()): |
| 1028 | start = int(query_start_loc_np[idx]) |
| 1029 | end = int(query_start_loc_np[idx + 1]) |
| 1030 | canvas = slice(start, end) |
| 1031 | soft = sc_embeds[slot, : end - start] |
| 1032 | inputs_embeds[canvas] = self.model.self_conditioning( |
| 1033 | inputs_embeds[canvas], soft.to(inputs_embeds.dtype) |
| 1034 | ) |
| 1035 | |
| 1036 | def prepare_inputs(self, input_batch, req_states) -> dict[str, Any]: |
| 1037 | states = self.diffusion_states |
| 1038 | num_tokens = input_batch.num_tokens |
| 1039 | num_reqs = input_batch.num_reqs |
| 1040 | |
| 1041 | # Write into the PERSISTENT inputs_embeds buffer so FULL CUDA graph |
| 1042 | # replay sees the latest values at the captured address. |
| 1043 | num_tokens_padded = input_batch.num_tokens_after_padding |
| 1044 | inputs_embeds = self._inputs_embeds_buf[:num_tokens_padded] |
| 1045 | |
| 1046 | # Populate embeddings: merge MM features when available, |
| 1047 | # otherwise embed input_ids as text-only. |
| 1048 | input_ids = input_batch.input_ids[:num_tokens] |
| 1049 | if self._pending_mm_embeds is not None: |
| 1050 | mm_embeds, is_mm_embed = self._pending_mm_embeds |
| 1051 | self._pending_mm_embeds = None |
| 1052 | inputs_embeds[:num_tokens].copy_( |
| 1053 | self.model.embed_input_ids( |
| 1054 | input_ids, |
| 1055 | multimodal_embeddings=mm_embeds, |
| 1056 | is_multimodal=is_mm_embed, |
| 1057 | ) |
| 1058 | ) |
| 1059 | else: |
| 1060 | inputs_embeds[:num_tokens].copy_(self.model.embed_input_ids(input_ids)) |
| 1061 | |
| 1062 | # Apply self-conditioning ONLY for denoising decode requests. |
| 1063 | states.step_allowed = None |
| 1064 | if input_batch.num_draft_tokens > 0 and self._req_id_to_index: |
| 1065 | slots_np = input_batch.idx_mapping_np[:num_reqs] |
| 1066 | num_logits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1]) |
| 1067 | is_decode_indices_np = np.where(num_logits_np > 0)[0] |
| 1068 | states.step_allowed = states.batch_allowed( |
| 1069 | slots_np[is_decode_indices_np].tolist() |
| 1070 | ) |
| 1071 | self._apply_self_conditioning( |
| 1072 | slots_np[is_decode_indices_np], |
| 1073 | is_decode_indices_np, |
| 1074 | input_batch.query_start_loc_np, |
| 1075 | inputs_embeds, |
| 1076 | states.self_conditioning_embeds, |
| 1077 | ) |
| 1078 | |
| 1079 | return {"inputs_embeds": inputs_embeds} |
| 1080 | |
| 1081 | def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: |
| 1082 | # CUDA graph capture path — return a slice of the SAME persistent |
| 1083 | # inputs_embeds buffer that `prepare_inputs` writes to at runtime, |
| 1084 | # so the captured graph and runtime point to identical addresses. |
| 1085 | return {"inputs_embeds": self._inputs_embeds_buf[:num_tokens]} |
| 1086 | |
| 1087 | def postprocess_state( |
| 1088 | self, idx_mapping, num_sampled, num_computed_tokens=None |
| 1089 | ) -> None: |
| 1090 | return None |
| 1091 | |
| 1092 | def prepare_attn( |
| 1093 | self, |
| 1094 | input_batch, |
| 1095 | cudagraph_mode, |
| 1096 | block_tables, |
| 1097 | slot_mappings, |
| 1098 | attn_groups, |
| 1099 | kv_cache_config, |
| 1100 | for_capture=False, |
| 1101 | ubatch_idx: int = 0, |
| 1102 | ) -> dict[str, Any]: |
| 1103 | assert ubatch_idx == 0, "DBO is not supported" |
| 1104 | if cudagraph_mode == CUDAGraphMode.FULL: |
| 1105 | num_reqs = input_batch.num_reqs_after_padding |
| 1106 | num_tokens = input_batch.num_tokens_after_padding |
| 1107 | else: |
| 1108 | num_reqs = input_batch.num_reqs |
| 1109 | num_tokens = input_batch.num_tokens |
| 1110 | |
| 1111 | query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) |
| 1112 | max_query_len = input_batch.num_scheduled_tokens.max().item() |
| 1113 | |
| 1114 | # Per-request causal mode: encoder (commit) = causal, |
| 1115 | # denoise = bidirectional. Pass GPU tensor so the attention |
| 1116 | # backend can handle mixed batches. |
| 1117 | actual_num_reqs = input_batch.num_reqs |
| 1118 | slots = input_batch.idx_mapping[:actual_num_reqs] |
| 1119 | # Invariant: the sampler flips is_encoder_phase to False only after a |
| 1120 | # request's FINAL prompt chunk, so a prompt spanning multiple chunks |
| 1121 | # (longer than the token budget) stays causal for every chunk. |
| 1122 | self._causal_buf[:actual_num_reqs].copy_( |
| 1123 | self.diffusion_states.is_encoder_phase[slots] |
| 1124 | ) |
| 1125 | if actual_num_reqs < num_reqs: |
| 1126 | self._causal_buf[actual_num_reqs:num_reqs] = 0 |
| 1127 | causal: bool | torch.Tensor = self._causal_buf[:num_reqs] |
| 1128 | |
| 1129 | return build_attn_metadata( |
| 1130 | attn_groups=attn_groups, |
| 1131 | num_reqs=num_reqs, |
| 1132 | num_tokens=num_tokens, |
| 1133 | query_start_loc_gpu=input_batch.query_start_loc, |
| 1134 | query_start_loc_cpu=query_start_loc_cpu, |
| 1135 | max_query_len=max_query_len, |
| 1136 | seq_lens=input_batch.seq_lens, |
| 1137 | max_seq_len=self.max_model_len, |
| 1138 | block_tables=block_tables, |
| 1139 | slot_mappings=slot_mappings, |
| 1140 | kv_cache_config=kv_cache_config, |
| 1141 | causal=causal, |
| 1142 | ) |
| 1143 | |
| 1144 | num_new_sampled_tokens_per_step: int = 0 |
| 1145 | |
| 1146 | |
| 1147 | # Penalty stub for the diffusion path: the runner reads |
| 1148 | # penalties_state.output_bin_counts, and post_update treats None as |
| 1149 | # "no penalty bookkeeping". |
| 1150 | _NO_PENALTIES_STATE = SimpleNamespace(output_bin_counts=None) |
| 1151 | |
| 1152 | |
| 1153 | class DiffusionSampler: |
| 1154 | """Batched accept/renoise sampler for DiffusionGemma. |
| 1155 | |
| 1156 | Follows the same structure as ``vllm.v1.worker.gpu.sample.sampler.Sampler``: |
| 1157 | decomposed into named methods, all GPU state in pre-allocated buffers, |
| 1158 | no GPU→CPU syncs on the hot path. |
| 1159 | """ |
| 1160 | |
| 1161 | def __init__( |
| 1162 | self, |
| 1163 | sampler: Any, |
| 1164 | diffusion_config: Any, |
| 1165 | vocab_size: int, |
| 1166 | diffusion_states: DiffusionGemmaRequestStates, |
| 1167 | *, |
| 1168 | confidence_threshold: float, |
| 1169 | t_min: float, |
| 1170 | t_max: float, |
| 1171 | entropy_bound: float, |
| 1172 | embed_weight: torch.Tensor, |
| 1173 | normalizer: torch.Tensor, |
| 1174 | sc_vocab_start: int = 0, |
| 1175 | sc_vocab_end: int | None = None, |
| 1176 | tp_size: int = 1, |
| 1177 | tp_group_name: str = "", |
| 1178 | ): |
| 1179 | self.sampling_states = sampler.sampling_states |
| 1180 | self.logprob_token_ids_state = sampler.logprob_token_ids_state |
| 1181 | self.req_states = sampler.req_states |
| 1182 | self.logits_mode = sampler.logprobs_mode in ("raw_logits", "processed_logits") |
| 1183 | # Self-conditioning soft embed = probs @ embed_weight * normalizer, |
| 1184 | # computed in the sampler (see _compiled_sample_step). ``embed_weight`` |
| 1185 | # is the vocab-parallel shard; [sc_vocab_start, sc_vocab_end) is this |
| 1186 | # rank's slice of the full vocab and tp_* drive the cross-rank |
| 1187 | # all-reduce. |
| 1188 | self.embed_weight = embed_weight |
| 1189 | self.normalizer = normalizer |
| 1190 | self.sc_vocab_start = sc_vocab_start |
| 1191 | self.sc_vocab_end = sc_vocab_end if sc_vocab_end is not None else vocab_size |
| 1192 | self.tp_size = tp_size |
| 1193 | self.tp_group_name = tp_group_name |
| 1194 | self.canvas_length = ( |
| 1195 | diffusion_config.canvas_length if diffusion_config is not None else 32 |
| 1196 | ) |
| 1197 | self.t_min = t_min |
| 1198 | self.t_max = t_max |
| 1199 | self.confidence_threshold = confidence_threshold |
| 1200 | self.vocab_size = vocab_size |
| 1201 | self.diffusion_states = diffusion_states |
| 1202 | self.entropy_bound = entropy_bound |
| 1203 | |
| 1204 | max_num_reqs = diffusion_states.max_num_reqs |
| 1205 | device = diffusion_states.device |
| 1206 | self._sampled = torch.zeros( |
| 1207 | max_num_reqs, |
| 1208 | self.canvas_length, |
| 1209 | dtype=torch.int32, |
| 1210 | device=device, |
| 1211 | ) |
| 1212 | self._num_sampled = torch.zeros( |
| 1213 | max_num_reqs, |
| 1214 | dtype=torch.int32, |
| 1215 | device=device, |
| 1216 | ) |
| 1217 | self._decode_slots = UvaBackedTensor(max_num_reqs, dtype=torch.int64) |
| 1218 | self._decode_idx = UvaBackedTensor(max_num_reqs, dtype=torch.int64) |
| 1219 | self._query_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32) |
| 1220 | self._num_logits = UvaBackedTensor(max_num_reqs, dtype=torch.int32) |
| 1221 | |
| 1222 | # Per-slot stash for logprobs computed on the converging denoise step. |
| 1223 | # Populated after the post-sample kernel detects convergence; consumed |
| 1224 | # on the subsequent commit step when num_sampled=CANVAS_LEN. |
| 1225 | self._pending_logprobs: dict[int, LogprobsTensors] = {} |
| 1226 | |
| 1227 | def add_request(self, req_idx: int, sampling_params: Any) -> None: |
| 1228 | if use_penalty(sampling_params): |
| 1229 | logger.warning_once( |
| 1230 | "DiffusionGemma does not support repetition/frequency/presence " |
| 1231 | "penalties; ignoring them for this request." |
| 1232 | ) |
| 1233 | # Purge any stale logprobs stashed under this slot by a prior request |
| 1234 | # that was aborted between its converging denoise and commit steps. |
| 1235 | self._pending_logprobs.pop(req_idx, None) |
| 1236 | self.sampling_states.add_request(req_idx, sampling_params) |
| 1237 | self.logprob_token_ids_state.add_request(req_idx, sampling_params) |
| 1238 | extra = getattr(sampling_params, "extra_args", None) or {} |
| 1239 | states = self.diffusion_states |
| 1240 | cap = extra.get("diffusion_max_steps") |
| 1241 | if cap is not None: |
| 1242 | cap = max(1, min(int(cap), states.max_denoising_steps)) |
| 1243 | states.max_steps[req_idx].fill_(cap) |
| 1244 | if cap == 1: |
| 1245 | states.single_step_slots.add(req_idx) |
| 1246 | width = extra.get("diffusion_canvas_length") |
| 1247 | if width: |
| 1248 | states.canvas_width_np[req_idx] = max( |
| 1249 | 1, min(int(width), self.canvas_length) |
| 1250 | ) |
| 1251 | width = int(states.canvas_width_np[req_idx]) |
| 1252 | seed = extra.get("diffusion_seed_canvas") |
| 1253 | if seed is not None: |
| 1254 | if len(seed) != width: |
| 1255 | raise ValueError( |
| 1256 | f"diffusion_seed_canvas must hold exactly {width} ids, " |
| 1257 | f"got {len(seed)}" |
| 1258 | ) |
| 1259 | states.set_seed_canvas(req_idx, seed) |
| 1260 | pins = extra.get("diffusion_pinned") |
| 1261 | if pins and seed is not None: |
| 1262 | states.set_pins(req_idx, [int(p) for p in pins]) |
| 1263 | if extra.get("diffusion_read_only"): |
| 1264 | states.set_read_only(req_idx) |
| 1265 | if extra.get("diffusion_constrained"): |
| 1266 | ids = list(getattr(sampling_params, "logprob_token_ids", None) or []) |
| 1267 | if not ids: |
| 1268 | raise ValueError( |
| 1269 | "diffusion_constrained needs logprob_token_ids: they are " |
| 1270 | "the allowed set." |
| 1271 | ) |
| 1272 | if self.tp_size > 1: |
| 1273 | raise ValueError("diffusion_constrained needs tensor parallel 1.") |
| 1274 | states.constrained[req_idx] = tuple(int(t) for t in ids) |
| 1275 | |
| 1276 | def apply_staged_writes(self) -> None: |
| 1277 | self.sampling_states.apply_staged_writes() |
| 1278 | self.logprob_token_ids_state.apply_staged_writes() |
| 1279 | |
| 1280 | @property |
| 1281 | def penalties_state(self): |
| 1282 | # Diffusion applies no penalties. The runner reads |
| 1283 | # penalties_state.output_bin_counts, so expose a stub holding None; |
| 1284 | # post_update treats None bin counts as "no penalty bookkeeping". |
| 1285 | return _NO_PENALTIES_STATE |
| 1286 | |
| 1287 | # ------------------------------------------------------------------ |
| 1288 | # Prefill |
| 1289 | # ------------------------------------------------------------------ |
| 1290 | |
| 1291 | def _finish_prefills( |
| 1292 | self, input_batch: Any, prefill_indices_np: np.ndarray |
| 1293 | ) -> None: |
| 1294 | """Transition requests whose prompt completes this step to denoising. |
| 1295 | |
| 1296 | Initializes their canvas, seeds draft tokens, and flips |
| 1297 | is_encoder_phase to False. Mid-chunk requests (prompt longer than the |
| 1298 | token budget) are left untouched so is_encoder_phase stays True and |
| 1299 | prepare_attn keeps causal attention for their remaining chunks. |
| 1300 | """ |
| 1301 | states = self.diffusion_states |
| 1302 | done_prefill_np = ( |
| 1303 | input_batch.num_computed_prefill_tokens_np[prefill_indices_np] |
| 1304 | + input_batch.num_scheduled_tokens[prefill_indices_np] |
| 1305 | >= input_batch.prefill_len_np[prefill_indices_np] |
| 1306 | ) |
| 1307 | ps = input_batch.idx_mapping_np[prefill_indices_np[done_prefill_np]] |
| 1308 | if len(ps) == 0: |
| 1309 | return |
| 1310 | # Move the slot indices across once, up front: indexing a device |
| 1311 | # tensor with a numpy array copies them over synchronously each time. |
| 1312 | ps_gpu = async_tensor_h2d( |
| 1313 | ps.astype(np.int64), device=states.is_encoder_phase.device |
| 1314 | ) |
| 1315 | states.init_canvas(ps_gpu) |
| 1316 | states.apply_seed_canvases(ps, ps_gpu) |
| 1317 | self.req_states.draft_tokens[ps_gpu, : self.canvas_length] = states.canvas[ |
| 1318 | ps_gpu |
| 1319 | ] |
| 1320 | states.is_encoder_phase.index_fill_(0, ps_gpu, False) |
| 1321 | |
| 1322 | def _handle_prefill( |
| 1323 | self, |
| 1324 | input_batch: Any, |
| 1325 | device: torch.device, |
| 1326 | ) -> SamplerOutput: |
| 1327 | num_reqs = input_batch.num_reqs |
| 1328 | self._finish_prefills(input_batch, np.arange(num_reqs)) |
| 1329 | sampled = self._sampled[:num_reqs, :1] |
| 1330 | sampled.zero_() |
| 1331 | num_sampled = self._num_sampled[:num_reqs] |
| 1332 | num_sampled.zero_() |
| 1333 | return SamplerOutput( |
| 1334 | sampled_token_ids=sampled, |
| 1335 | logprobs_tensors=None, |
| 1336 | num_nans=None, |
| 1337 | num_sampled=num_sampled, |
| 1338 | num_rejected=num_sampled, |
| 1339 | ) |
| 1340 | |
| 1341 | # ------------------------------------------------------------------ |
| 1342 | # Decode helpers |
| 1343 | # ------------------------------------------------------------------ |
| 1344 | |
| 1345 | def _build_output( |
| 1346 | self, |
| 1347 | input_batch: Any, |
| 1348 | sampled: torch.Tensor, |
| 1349 | num_sampled: torch.Tensor, |
| 1350 | per_req_nlogits_np: np.ndarray, |
| 1351 | device: torch.device, |
| 1352 | logprobs_tensors: LogprobsTensors | None = None, |
| 1353 | ) -> SamplerOutput: |
| 1354 | """Compute num_rejected and build SamplerOutput.""" |
| 1355 | num_reqs = input_batch.num_reqs |
| 1356 | |
| 1357 | self._query_lens.np[:num_reqs] = np.diff( |
| 1358 | input_batch.query_start_loc_np[: num_reqs + 1] |
| 1359 | ) |
| 1360 | self._num_logits.np[:num_reqs] = per_req_nlogits_np |
| 1361 | self._query_lens.copy_to_uva() |
| 1362 | self._num_logits.copy_to_uva() |
| 1363 | |
| 1364 | num_rejected = _compute_num_rejected( |
| 1365 | self._num_logits.gpu[:num_reqs], |
| 1366 | num_sampled, |
| 1367 | input_batch.query_start_loc[: num_reqs + 1], |
| 1368 | ) |
| 1369 | |
| 1370 | return SamplerOutput( |
| 1371 | sampled_token_ids=sampled, |
| 1372 | logprobs_tensors=logprobs_tensors, |
| 1373 | num_nans=None, |
| 1374 | num_sampled=num_sampled, |
| 1375 | num_rejected=num_rejected, |
| 1376 | ) |
| 1377 | |
| 1378 | # ------------------------------------------------------------------ |
| 1379 | # Main entry point |
| 1380 | # ------------------------------------------------------------------ |
| 1381 | |
| 1382 | def __call__( |
| 1383 | self, |
| 1384 | logits: torch.Tensor, |
| 1385 | input_batch: Any, |
| 1386 | draft_logits: torch.Tensor | None = None, |
| 1387 | ) -> SamplerOutput: |
| 1388 | num_reqs = input_batch.num_reqs |
| 1389 | device = logits.device |
| 1390 | |
| 1391 | if input_batch.num_draft_tokens == 0: |
| 1392 | return self._handle_prefill(input_batch, device) |
| 1393 | |
| 1394 | # --- CPU/NumPy setup (outside compile): split decode vs prefill, init |
| 1395 | # canvas for any new prefills, and stage decode slot indices to GPU. --- |
| 1396 | states = self.diffusion_states |
| 1397 | slots_np = input_batch.idx_mapping_np[:num_reqs] |
| 1398 | per_req_nlogits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1]) |
| 1399 | |
| 1400 | decode_indices_np = np.where(per_req_nlogits_np > 0)[0] |
| 1401 | prefill_indices_np = np.where(per_req_nlogits_np == 0)[0] |
| 1402 | decode_slots_np = slots_np[decode_indices_np] |
| 1403 | |
| 1404 | if len(prefill_indices_np) > 0: |
| 1405 | self._finish_prefills(input_batch, prefill_indices_np) |
| 1406 | |
| 1407 | num_decode = len(decode_indices_np) |
| 1408 | self._decode_slots.np[:num_decode] = decode_slots_np |
| 1409 | self._decode_idx.np[:num_decode] = decode_indices_np |
| 1410 | self._decode_slots.copy_to_uva() |
| 1411 | self._decode_idx.copy_to_uva() |
| 1412 | decode_slots = self._decode_slots.gpu[:num_decode] |
| 1413 | decode_idx = self._decode_idx.gpu[:num_decode] |
| 1414 | |
| 1415 | # Real canvas length per decode request. Equals CL except when a canvas |
| 1416 | # was truncated near max_model_len, in which case the scheduler gave us |
| 1417 | # fewer than CL logits for that request. |
| 1418 | valid_canvas_len_np = per_req_nlogits_np[per_req_nlogits_np > 0] |
| 1419 | valid_canvas_len = async_tensor_h2d( |
| 1420 | valid_canvas_len_np.astype(np.int64), device=device |
| 1421 | ) |
| 1422 | |
| 1423 | # Per-request top_k/top_p, mirroring the AR sampler. Masked tokens |
| 1424 | # become -inf and survive the temperature scaling in the compiled |
| 1425 | # step, so Gumbel sampling, probs, and entropy all see the filtered |
| 1426 | # distribution. The committed argmax (always the top-1 token) is |
| 1427 | # unaffected; only the canvas exploration is constrained. Applied |
| 1428 | # before canvas padding so phantom positions stay uniform. |
| 1429 | if num_decode > 0: |
| 1430 | top_k, top_p = self.sampling_states.get_top_k_top_p( |
| 1431 | decode_slots.repeat_interleave( |
| 1432 | valid_canvas_len, output_size=int(valid_canvas_len_np.sum()) |
| 1433 | ), |
| 1434 | decode_slots_np, |
| 1435 | ) |
| 1436 | if top_k is not None or top_p is not None: |
| 1437 | logits = apply_top_k_top_p(logits.float(), top_k, top_p) |
| 1438 | |
| 1439 | # Where each decode request's rows start in the flat logits. Tiles |
| 1440 | # below gather and pad a request to the tile's width. |
| 1441 | row_starts_np = np.concatenate( |
| 1442 | ([0], np.cumsum(valid_canvas_len_np)[:-1]) |
| 1443 | ).astype(np.int64) |
| 1444 | row_starts = async_tensor_h2d(row_starts_np, device=device) |
| 1445 | |
| 1446 | # Clear once: the tiled loop below only scatters its own decode slots, |
| 1447 | # so it must not re-clear earlier tiles' writes. |
| 1448 | sampled = self._sampled[:num_reqs] |
| 1449 | num_sampled = self._num_sampled[:num_reqs] |
| 1450 | sampled.zero_() |
| 1451 | num_sampled.zero_() |
| 1452 | |
| 1453 | all_slots = input_batch.idx_mapping[:num_reqs] |
| 1454 | |
| 1455 | # Snapshot which slots are committing BEFORE the compiled step runs, |
| 1456 | # since it mutates is_encoder_phase (commit→False, converge→True). |
| 1457 | is_committing = states.is_encoder_phase[decode_slots].clone() |
| 1458 | |
| 1459 | # Constrained step: logits are [rows, K] over the shared allowed set, |
| 1460 | # so the self-conditioning matmul only needs those K embedding rows. |
| 1461 | allowed = states.step_allowed |
| 1462 | embed_weight = self.embed_weight |
| 1463 | vocab_size = self.vocab_size |
| 1464 | sc_vocab_start, sc_vocab_end = self.sc_vocab_start, self.sc_vocab_end |
| 1465 | if allowed is not None: |
| 1466 | assert logits.shape[-1] == allowed.numel() |
| 1467 | embed_weight = self.embed_weight[allowed] |
| 1468 | sc_vocab_start, sc_vocab_end = 0, allowed.numel() |
| 1469 | elif states.constrained and num_decode > 0: |
| 1470 | # The decode slots do not share one set, so this step runs over |
| 1471 | # the full vocabulary. Mask each constrained request's rows to |
| 1472 | # its own set so it reads the same as on the shared path. The |
| 1473 | # helper masks a copy, since the runner owns `logits`: one |
| 1474 | # full-vocab copy per step, as for top_k/top_p above. |
| 1475 | per_row = [ |
| 1476 | None if ids is None else states.allowed_tensor(ids) |
| 1477 | for ids in (states.constrained.get(s) for s in decode_slots_np.tolist()) |
| 1478 | ] |
| 1479 | logits = _mask_rows_to_allowed( |
| 1480 | logits, row_starts_np, valid_canvas_len_np, per_row |
| 1481 | ) |
| 1482 | |
| 1483 | slots_np = input_batch.idx_mapping_np[:num_reqs] |
| 1484 | max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np) |
| 1485 | # Requests may ask for specific token ids' logprobs instead of, or as |
| 1486 | # well as, a top-k. |
| 1487 | max_token_ids = self.logprob_token_ids_state.max_num_token_ids(slots_np) |
| 1488 | want_logprobs = max_num_logprobs >= 0 or max_token_ids > 0 |
| 1489 | num_logprobs = max(max_num_logprobs, 0) |
| 1490 | |
| 1491 | # Sample per tile. Decode requests are grouped by canvas width and each |
| 1492 | # tile runs the compiled step at that width over [:, :W] views of the |
| 1493 | # state, so a narrow read pays for its own rows rather than the served |
| 1494 | # canvas. Widths ascend, so the last tile's canvas-to-draft copy (over |
| 1495 | # all slots) is the widest. The fp32 pipeline keeps several live |
| 1496 | # [tile * W, vocab] copies, so a tile is also bounded by free memory. |
| 1497 | widths_np = states.canvas_width_np[decode_slots_np] |
| 1498 | order = np.argsort(widths_np, kind="stable") |
| 1499 | free = torch.accelerator.get_memory_info()[0] if num_decode > 0 else 0 |
| 1500 | run_start = 0 |
| 1501 | while run_start < num_decode: |
| 1502 | W = int(widths_np[order[run_start]]) |
| 1503 | run_end = run_start |
| 1504 | while run_end < num_decode and widths_np[order[run_end]] == W: |
| 1505 | run_end += 1 |
| 1506 | # Transient [tile * W, vocab] tensors: the softmax in the embedding |
| 1507 | # dtype for self-conditioning, plus the fp32 scaled logits when |
| 1508 | # logprobs are wanted (pad for allocator overhead). |
| 1509 | budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 3, 1)) |
| 1510 | for t0 in range(run_start, run_end, budget): |
| 1511 | sel_np = order[t0 : min(t0 + budget, run_end)] |
| 1512 | n = len(sel_np) |
| 1513 | sel = async_tensor_h2d(sel_np.astype(np.int64), device=device) |
| 1514 | tile_slots = decode_slots[sel] |
| 1515 | tile_valid = valid_canvas_len[sel] |
| 1516 | tile_valid_np = valid_canvas_len_np[sel_np] |
| 1517 | contiguous = bool((sel_np == np.arange(sel_np[0], sel_np[0] + n)).all()) |
| 1518 | if contiguous and tile_valid_np.min() == W: |
| 1519 | r0 = int(row_starts_np[sel_np[0]]) |
| 1520 | tile_logits = logits[r0 : r0 + n * W] |
| 1521 | else: |
| 1522 | # Pad each request to W. Phantom positions are zeroed: |
| 1523 | # uniform logits, high entropy, argmax 0, never committed. |
| 1524 | # masked_fill, not multiply, so -inf from top_k/top_p |
| 1525 | # filtering above cannot turn a phantom row into NaN. |
| 1526 | ar = torch.arange(W, device=device) |
| 1527 | src = (row_starts[sel].unsqueeze(1) + ar.unsqueeze(0)).clamp_max( |
| 1528 | logits.shape[0] - 1 |
| 1529 | ) |
| 1530 | valid = ar.unsqueeze(0) < tile_valid.unsqueeze(1) |
| 1531 | tile_logits = logits[src.reshape(-1)].masked_fill_( |
| 1532 | ~valid.reshape(-1, 1), 0 |
| 1533 | ) |
| 1534 | compute_sc = ( |
| 1535 | not states.single_step_slots |
| 1536 | or not states.single_step_slots.issuperset( |
| 1537 | decode_slots_np[sel_np].tolist() |
| 1538 | ) |
| 1539 | ) |
| 1540 | |
| 1541 | temp = _denoise_temperature( |
| 1542 | states.step, |
| 1543 | tile_slots, |
| 1544 | float(states.max_denoising_steps), |
| 1545 | self.t_min, |
| 1546 | self.t_max, |
| 1547 | ) |
| 1548 | probs_dtype = self.embed_weight.dtype if compute_sc else None |
| 1549 | if tile_logits.is_cuda: |
| 1550 | seed = int(torch.randint(0, 2**31 - 1, (1,)).item()) |
| 1551 | stats = sample_row_stats(tile_logits, temp, W, seed, probs_dtype) |
| 1552 | else: |
| 1553 | stats = sample_row_stats_reference( |
| 1554 | tile_logits, temp, W, probs_dtype |
| 1555 | ) |
| 1556 | argmax_rows, sample_rows, entropy_rows, probs_rows = stats |
| 1557 | if allowed is not None: |
| 1558 | # K-space picks back to token ids. The softmax stays K wide |
| 1559 | # for the K-row self-conditioning matmul. |
| 1560 | argmax_rows = allowed[argmax_rows] |
| 1561 | sample_rows = allowed[sample_rows] |
| 1562 | probs = probs_rows.view(n, W, -1) if probs_rows is not None else None |
| 1563 | # Only the logprob stash reads the scaled logits. |
| 1564 | scaled = None |
| 1565 | if want_logprobs: |
| 1566 | scaled = tile_logits.float().view(n, W, -1) / temp[ |
| 1567 | :, None, None |
| 1568 | ].clamp(min=1e-10) |
| 1569 | |
| 1570 | _compiled_sample_step( |
| 1571 | sample_rows.view(n, W), |
| 1572 | argmax_rows.view(n, W), |
| 1573 | entropy_rows.view(n, W), |
| 1574 | probs, |
| 1575 | tile_slots, |
| 1576 | decode_idx[sel], |
| 1577 | all_slots, |
| 1578 | tile_valid, |
| 1579 | # State, viewed at this tile's width |
| 1580 | states.canvas[:, :W], |
| 1581 | states.argmax_canvas[:, :W], |
| 1582 | states.step, |
| 1583 | states.is_encoder_phase, |
| 1584 | states.confident, |
| 1585 | states.self_conditioning_embeds[:, :W], |
| 1586 | embed_weight, |
| 1587 | self.normalizer, |
| 1588 | states.accepted_canvas_history[:, :, :W], |
| 1589 | states.accepted_canvas_history_len, |
| 1590 | states.max_steps, |
| 1591 | states.pin_mask[:, :W], |
| 1592 | states.seed_canvas[:, :W], |
| 1593 | states.read_only, |
| 1594 | # Output |
| 1595 | sampled[:, :W], |
| 1596 | num_sampled, |
| 1597 | self.req_states.draft_tokens, |
| 1598 | # Config |
| 1599 | confidence_threshold=self.confidence_threshold, |
| 1600 | vocab_size=vocab_size, |
| 1601 | CL=W, |
| 1602 | ST=states.stability_threshold, |
| 1603 | entropy_bound=self.entropy_bound, |
| 1604 | sc_vocab_start=sc_vocab_start, |
| 1605 | sc_vocab_end=sc_vocab_end, |
| 1606 | tp_size=self.tp_size, |
| 1607 | tp_group_name=self.tp_group_name, |
| 1608 | compute_sc=compute_sc, |
| 1609 | ) |
| 1610 | |
| 1611 | # Stash newly converged logprobs, including reads that emit now. |
| 1612 | if want_logprobs: |
| 1613 | converged_mask = states.is_encoder_phase[tile_slots] | ( |
| 1614 | num_sampled[decode_idx[sel]] > 0 |
| 1615 | ) |
| 1616 | just_converged = converged_mask & ~is_committing[sel] |
| 1617 | if just_converged.any(): |
| 1618 | assert scaled is not None |
| 1619 | flat_logits = scaled.reshape(-1, scaled.shape[-1]) |
| 1620 | argmax_tokens = scaled.argmax(dim=-1) |
| 1621 | raw_flat: torch.Tensor | None = None |
| 1622 | for local_idx in just_converged.nonzero(as_tuple=True)[0]: |
| 1623 | li = local_idx.item() |
| 1624 | slot = tile_slots[local_idx].item() |
| 1625 | # Stash only the real canvas positions; padded tail |
| 1626 | # positions are never emitted. |
| 1627 | k_i = int(tile_valid_np[li]) |
| 1628 | pos = li * W |
| 1629 | src = flat_logits |
| 1630 | if slot in states.read_only_slots: |
| 1631 | # Read-only slots report logprobs at temperature 1. |
| 1632 | # The schedule-tempered logits share the argmax. |
| 1633 | if raw_flat is None: |
| 1634 | raw_flat = tile_logits.float() |
| 1635 | src = raw_flat |
| 1636 | if allowed is not None: |
| 1637 | # Column j of the K-space logits is allowed[j]: |
| 1638 | # report every allowed id, renormalized over |
| 1639 | # the set, with the argmax in column 0. |
| 1640 | rows = src[pos : pos + k_i] |
| 1641 | lp = rows.log_softmax(dim=-1) |
| 1642 | am = argmax_tokens[local_idx][:k_i] |
| 1643 | sel_lp = lp.gather(1, am.unsqueeze(1)) |
| 1644 | # Same convention as _ranks_kernel: 1-based, |
| 1645 | # ties count, so the rank is the number of |
| 1646 | # columns at or above the selected one. Out |
| 1647 | # of set ids have zero probability, so the |
| 1648 | # rank within the set is the real rank. |
| 1649 | ranks = (lp >= sel_lp).sum(dim=1) |
| 1650 | self._pending_logprobs[slot] = LogprobsTensors( |
| 1651 | logprob_token_ids=torch.cat( |
| 1652 | ( |
| 1653 | allowed[am].unsqueeze(1), |
| 1654 | allowed.unsqueeze(0).expand(k_i, -1), |
| 1655 | ), |
| 1656 | dim=1, |
| 1657 | ), |
| 1658 | logprobs=torch.cat((sel_lp, lp), dim=1), |
| 1659 | selected_token_ranks=ranks.to(torch.int64), |
| 1660 | ) |
| 1661 | continue |
| 1662 | per_req_ids = max_token_ids > 0 |
| 1663 | self._pending_logprobs[slot] = compute_topk_scores( |
| 1664 | src[pos : pos + k_i], |
| 1665 | num_logprobs, |
| 1666 | argmax_tokens[local_idx][:k_i], |
| 1667 | logprob_token_ids_state=( |
| 1668 | self.logprob_token_ids_state |
| 1669 | if per_req_ids |
| 1670 | else None |
| 1671 | ), |
| 1672 | # every row of this stash belongs to one slot |
| 1673 | expanded_idx_mapping=( |
| 1674 | torch.full( |
| 1675 | (k_i,), slot, dtype=torch.int32, device=device |
| 1676 | ) |
| 1677 | if per_req_ids |
| 1678 | else None |
| 1679 | ), |
| 1680 | max_per_req_token_ids=max_token_ids, |
| 1681 | logits_mode=self.logits_mode, |
| 1682 | ) |
| 1683 | run_start = run_end |
| 1684 | |
| 1685 | # Only emitting requests consume their stashed logprobs. |
| 1686 | logprobs_tensors = None |
| 1687 | if want_logprobs and self._pending_logprobs: |
| 1688 | emitting_slots = set(slots_np[num_sampled.cpu().numpy() > 0].tolist()) |
| 1689 | parts: list[LogprobsTensors] = [] |
| 1690 | cu_gen: list[int] = [] |
| 1691 | flat_offset = 0 |
| 1692 | for i in range(num_reqs): |
| 1693 | cu_gen.append(flat_offset) |
| 1694 | slot = int(slots_np[i]) |
| 1695 | if slot in emitting_slots and slot in self._pending_logprobs: |
| 1696 | lp = self._pending_logprobs.pop(slot) |
| 1697 | parts.append(lp) |
| 1698 | flat_offset += lp.logprobs.shape[0] |
| 1699 | if parts: |
| 1700 | logprobs_tensors = _concat_logprob_stashes(parts, cu_gen) |
| 1701 | |
| 1702 | return self._build_output( |
| 1703 | input_batch, |
| 1704 | sampled, |
| 1705 | num_sampled, |
| 1706 | per_req_nlogits_np, |
| 1707 | device, |
| 1708 | logprobs_tensors=logprobs_tensors, |
| 1709 | ) |