diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0b69ac8f4753..820f66d46681 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -446,6 +446,7 @@ th { | `MixtralForCausalLM` | Mixtral-8x7B, Mixtral-8x7B-Instruct | `mistralai/Mixtral-8x7B-v0.1`, `mistralai/Mixtral-8x7B-Instruct-v0.1`, `mistral-community/Mixtral-8x22B-v0.1`, etc. | ✅︎ | ✅︎ | | `NemotronForCausalLM` | Nemotron-3, Nemotron-4, Minitron | `nvidia/Minitron-8B-Base`, `mgoin/Nemotron-4-340B-Base-hf-FP8`, etc. | ✅︎ | ✅︎ | | `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`, `nvidia/Nemotron-H-47B-Base-8K`, `nvidia/Nemotron-H-56B-Base-8K`, etc. | ✅︎ | ✅︎ | +| `NemotronLabsDiffusionModel`, `NemotronLabsDiffusionForCausalLM` | Nemotron Labs Diffusion | `nvidia/Nemotron-Labs-Diffusion-3B`, `nvidia/Nemotron-Labs-Diffusion-8B` | | | | `OlmoHybridForCausalLM` | OLMo Hybrid | `allenai/Olmo-Hybrid-7B` | ✅︎ | ✅︎ | | `OlmoeForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ | | `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ | @@ -502,6 +503,9 @@ Some models are supported only via the [Transformers modeling backend](#transfor !!! note Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096. +!!! note + Nemotron Labs Diffusion runs masked block diffusion by default. `--diffusion-config` sets `canvas_length` (default: the checkpoint's `block_size`) and `model_options` such as `selection_policy` and `confidence_threshold`; request temperatures are applied as usual. Pass `--hf-overrides '{"ar_mode": true}'` (or select `NemotronLabsDiffusionForCausalLM`) for ordinary autoregressive generation with the same checkpoint. + ## List of Multimodal Language Models The following modalities are supported depending on the model: diff --git a/tests/models/language/generation/test_nemotron_dllm.py b/tests/models/language/generation/test_nemotron_dllm.py new file mode 100644 index 000000000000..70d4c43c533e --- /dev/null +++ b/tests/models/language/generation/test_nemotron_dllm.py @@ -0,0 +1,873 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reveal-rule unit tests and GPU model parity/generation checks. + +Weight-dependent tests load nvidia/Nemotron-Labs-Diffusion-3B from the Hub by +default; set NEMOTRON_DLM_MODEL_PATH to another Hub ID or a local 3B or 8B +checkpoint directory. Synthetic reveal rules run eagerly on CPU; GPU +model tests exercise the compiled sampler with the checkpoint's vocabulary. +""" + +import json +import os +from functools import partial + +import pytest +import torch + +MODEL_PATH = os.environ.get( + "NEMOTRON_DLM_MODEL_PATH", + "nvidia/Nemotron-Labs-Diffusion-3B", +) +MASK_TOKEN_ID = 100 +CANVAS_LENGTH = 32 +MAX_MODEL_LEN = 2048 + + +def _weights_available() -> bool: + # A Hub ID is downloaded on first use; a local path must be a checkpoint. + if os.path.isdir(MODEL_PATH) or MODEL_PATH.startswith(("/", ".", "~")): + return os.path.isfile(os.path.join(MODEL_PATH, "config.json")) + return True + + +requires_weights = pytest.mark.skipif( + not _weights_available(), + reason=f"Nemotron-DLM weights not found at {MODEL_PATH}", +) +requires_gpu = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires a CUDA device" +) + + +# --------------------------------------------------------------------------- +# CPU-only: registration & diffusion detection +# --------------------------------------------------------------------------- +def test_config_registration(): + from vllm.transformers_utils.configs import NemotronLabsDiffusionConfig + + cfg = NemotronLabsDiffusionConfig(block_size=32, mask_token_id=100) + # canvas_length is the field ModelConfig.is_diffusion keys off of. + assert cfg.canvas_length == 32 + assert cfg.mask_token_id == 100 + + +@pytest.mark.parametrize( + "hidden_size,intermediate_size,num_layers,rope_factor", + [(3072, 9216, 26, 16.0), (4096, 14336, 34, 8.0)], + ids=["3b", "8b"], +) +def test_checkpoint_dimensions_and_rope( + tmp_path, hidden_size, intermediate_size, num_layers, rope_factor +): + """Loading 8B must preserve its dimensions and RoPE, not 3B defaults.""" + from vllm.transformers_utils.config import get_config + from vllm.transformers_utils.configs import NemotronLabsDiffusionConfig + + checkpoint = { + "model_type": "nemotron_labs_diffusion", + "architectures": ["NemotronLabsDiffusionModel"], + "hidden_size": hidden_size, + "intermediate_size": intermediate_size, + "num_hidden_layers": num_layers, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "head_dim": 128, + "max_position_embeddings": int(16384 * rope_factor), + "block_size": 32, + "mask_token_id": 100, + "rope_parameters": { + "rope_type": "yarn", + "rope_theta": 1000000.0, + "factor": rope_factor, + "original_max_position_embeddings": 16384, + "llama_4_scaling_beta": 0.1, + "mscale": 1.0, + "mscale_all_dim": 1.0, + }, + } + (tmp_path / "config.json").write_text(json.dumps(checkpoint)) + config = get_config(str(tmp_path), trust_remote_code=False) + assert isinstance(config, NemotronLabsDiffusionConfig) + for key, value in checkpoint.items(): + assert getattr(config, key) == value, key + assert config.canvas_length == checkpoint["block_size"] + + +def test_arch_in_model_registry(): + from vllm.model_executor.models.registry import ModelRegistry + + assert "NemotronLabsDiffusionModel" in ModelRegistry.get_supported_archs() + + +@pytest.mark.parametrize("flash_version", [None, 3]) +def test_attention_backend_supports_mixed_causality(flash_version): + from types import SimpleNamespace + from typing import cast + + from vllm.config import VllmConfig + from vllm.config.attention import AttentionConfig + from vllm.model_executor.models.config import ( + NemotronLabsDiffusionForBlockDiffusionConfig, + ) + from vllm.transformers_utils.configs import NemotronLabsDiffusionConfig + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + attention = AttentionConfig() + if flash_version is not None: + attention.backend = AttentionBackendEnum.FLASH_ATTN + attention.flash_attn_version = flash_version + config = cast( + VllmConfig, + SimpleNamespace( + attention_config=attention, + model_config=SimpleNamespace( + hf_config=NemotronLabsDiffusionConfig(), + override_generation_config={}, + is_diffusion=True, + ), + diffusion_config=None, + scheduler_config=None, + ), + ) + if flash_version is not None: + with pytest.raises(ValueError, match="requires FA4"): + NemotronLabsDiffusionForBlockDiffusionConfig.verify_and_update_config( + config + ) + else: + NemotronLabsDiffusionForBlockDiffusionConfig.verify_and_update_config(config) + assert attention.backend == AttentionBackendEnum.TRITON_ATTN + + +@pytest.mark.parametrize( + "architecture,ar_mode,expected", + [ + ("NemotronLabsDiffusionModel", False, True), + ("NemotronLabsDiffusionModel", True, False), + ("NemotronLabsDiffusionForCausalLM", False, False), + ("LLaDAModelLM", True, True), + ], +) +def test_ar_mode_diffusion_detection(architecture, ar_mode, expected): + from types import SimpleNamespace + from typing import cast + + from vllm.config import ModelConfig + from vllm.config.model_arch import ModelArchitectureConfig + from vllm.transformers_utils.configs import NemotronLabsDiffusionConfig + + # The Nemotron config drops its canvas for the AR selectors, so the + # generic canvas_length check in ModelConfig.is_diffusion applies. + # Other models keep a canvas regardless of an unrelated ar_mode flag. + if architecture.startswith("NemotronLabsDiffusion"): + hf_config = NemotronLabsDiffusionConfig( + block_size=32, architectures=[architecture] + ) + hf_config.ar_mode = ar_mode + assert (hf_config.canvas_length is not None) is expected + else: + hf_config = SimpleNamespace(canvas_length=32, ar_mode=ar_mode) + config = object.__new__(ModelConfig) + config.hf_config = hf_config + config.model_arch_config = cast( + ModelArchitectureConfig, SimpleNamespace(architectures=[architecture]) + ) + assert config.is_diffusion is expected + + +@pytest.mark.parametrize("explicit_diffusion_config", [False, True]) +def test_ar_mode_config(explicit_diffusion_config): + from types import SimpleNamespace + from typing import cast + + from vllm.config import VllmConfig + from vllm.config.attention import AttentionConfig + from vllm.config.diffusion import DiffusionConfig + from vllm.config.scheduler import SchedulerConfig + from vllm.model_executor.models.config import ( + NemotronLabsDiffusionForBlockDiffusionConfig, + ) + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + # AR permits FA3 and keeps ordinary scheduler concurrency defaults. + attention = AttentionConfig( + backend=AttentionBackendEnum.FLASH_ATTN, flash_attn_version=3 + ) + scheduler = SimpleNamespace(max_num_seqs=SchedulerConfig.DEFAULT_MAX_NUM_SEQS) + config = cast( + VllmConfig, + SimpleNamespace( + attention_config=attention, + model_config=SimpleNamespace( + is_diffusion=False, override_generation_config={} + ), + diffusion_config=( + DiffusionConfig(canvas_length=32) if explicit_diffusion_config else None + ), + scheduler_config=scheduler, + ), + ) + if explicit_diffusion_config: + with pytest.raises(ValueError, match="AR mode cannot be combined"): + NemotronLabsDiffusionForBlockDiffusionConfig.verify_and_update_config( + config + ) + else: + NemotronLabsDiffusionForBlockDiffusionConfig.verify_and_update_config(config) + assert config.diffusion_config is None + assert attention.flash_attn_version == 3 + assert scheduler.max_num_seqs == SchedulerConfig.DEFAULT_MAX_NUM_SEQS + assert config.model_config.override_generation_config["max_new_tokens"] is None + + +def test_diffusion_model_options_are_model_specific(): + """DiffusionConfig parses without a model; each model checks its own options.""" + from types import SimpleNamespace + from typing import Any, cast + + from vllm.config.diffusion import DiffusionConfig + from vllm.model_executor.models.config import ( + DiffusionGemmaDiffusionOptions, + NemotronDiffusionOptions, + _parse_diffusion_model_options, + ) + + def parse(options: dict[str, Any], cls: type) -> Any: + dc = DiffusionConfig(canvas_length=32, model_options=options) + config = cast(Any, SimpleNamespace(diffusion_config=dc)) + _parse_diffusion_model_options(config, cls) + return dc.model_options + + opts = parse({"temperature": 0.0}, NemotronDiffusionOptions) + assert opts.temperature == 0.0 + with pytest.raises(ValueError): + parse({"temperature": 0.0}, DiffusionGemmaDiffusionOptions) + with pytest.raises(ValueError): + parse({"temperature": -0.5}, NemotronDiffusionOptions) + + +# --------------------------------------------------------------------------- +# CPU, no weights: masked-diffusion step semantics +# --------------------------------------------------------------------------- +@pytest.fixture +def eager_sampler(): + """Test reveal rules independently of Inductor; model tests cover compilation.""" + with torch.compiler.set_stance("force_eager"): + yield + + +def test_masked_step_sampling_semantics(eager_sampler): + from vllm.model_executor.models.nemotron_dllm import _compiled_masked_step + + device = "cpu" + CL, V, MAX_REQS = 4, 12, 2 + mask_id = 10 + max_steps = 2 # even schedule: k = ceil(4/2) = 2 unmasked per step + + def fresh_state(): + return { + "canvas": torch.full( + (MAX_REQS, CL), mask_id, dtype=torch.int64, device=device + ), + "step": torch.zeros(MAX_REQS, dtype=torch.int32, device=device), + "phase": torch.zeros(MAX_REQS, dtype=torch.bool, device=device), + "lp": torch.zeros(MAX_REQS, CL, dtype=torch.float32, device=device), + "rk": torch.zeros(MAX_REQS, CL, dtype=torch.int32, device=device), + } + + def run_step(st, logits, temperature, leftmost=False, steps=max_steps): + sampled = torch.zeros(1, CL, dtype=torch.int32, device=device) + num_sampled = torch.zeros(1, dtype=torch.int32, device=device) + draft = torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device) + zero = torch.zeros(1, dtype=torch.int64, device=device) + _compiled_masked_step( + logits.reshape(-1, V), + zero, # decode_slots + zero, # decode_idx + zero, # all_slots + torch.full((1,), CL, dtype=torch.int64, device=device), + torch.full((1,), temperature, dtype=torch.float32, device=device), + st["canvas"], + st["step"], + st["phase"], + st["lp"], + st["rk"], + sampled, + num_sampled, + draft, + max_denoising_steps=steps, + mask_token_id=mask_id, + CL=CL, + all_greedy=temperature == 0, + capture_logprobs=True, + leftmost=leftmost, + threshold_mode=False, + log_threshold=0.0, + ) + + torch.manual_seed(0) + logits1 = torch.randn(1, CL, V, device=device) * 4.0 + logits1[..., mask_id] = -100.0 # never pick the mask token + + # Greedy (T=0): unmasked tokens == argmax, captured logprob/rank match + # the step's log_softmax and top-1 rank. + st = fresh_state() + run_step(st, logits1, 0.0) + lp1 = logits1[0].float().log_softmax(-1) + x0 = logits1[0].argmax(-1) + step1 = st["canvas"][0] != mask_id + assert int(step1.sum()) == 2 + assert torch.equal(st["canvas"][0][step1], x0[step1]) + want = lp1.gather(-1, x0.unsqueeze(-1)).squeeze(-1) + assert torch.allclose(st["lp"][0][step1], want[step1], atol=1e-5) + assert (st["rk"][0][step1] == 1).all() + + # At-unmask persistence: a second step with DIFFERENT logits must not + # rewrite the step-1 captures, and step-2 captures reflect step-2 logits. + lp_after1 = st["lp"][0].clone() + torch.manual_seed(1) + logits2 = torch.randn(1, CL, V, device=device) * 4.0 + logits2[..., mask_id] = -100.0 + run_step(st, logits2, 0.0) + assert not (st["canvas"][0] == mask_id).any() + assert torch.allclose(st["lp"][0][step1], lp_after1[step1]) + step2 = ~step1 + lp2 = logits2[0].float().log_softmax(-1) + want2 = lp2.gather(-1, st["canvas"][0].unsqueeze(-1)).squeeze(-1) + assert torch.allclose(st["lp"][0][step2], want2[step2], atol=1e-5) + + # Sampling (T=1): draws differ across seeds; the captured logprob/rank + # match the SAMPLED token under this step's distribution. + draws = [] + for seed in range(3): + torch.manual_seed(100 + seed) + st = fresh_state() + run_step(st, logits1 * 0.25, 1.0) # flatter -> diverse draws + unm = st["canvas"][0] != mask_id + toks = st["canvas"][0] + lp_ref = (logits1[0] * 0.25).float().log_softmax(-1) + got = lp_ref.gather(-1, toks.clamp(0, V - 1).unsqueeze(-1)).squeeze(-1) + assert torch.allclose(st["lp"][0][unm], got[unm], atol=1e-5) + ranks = (lp_ref >= got.unsqueeze(-1)).sum(-1).int() + assert torch.equal(st["rk"][0][unm], ranks[unm]) + draws.append(tuple(toks[unm].tolist())) + assert len(set(draws)) >= 2, f"3 seeds gave identical draws: {draws}" + + # Leftmost selection: reveal order is strictly position order even when + # a rightmost position is by far the most confident (leftmost-reveal RL + # rollouts need this to match the trainer's recompute factorization). + spiked = logits1.clone() + spiked[0, CL - 1, 5] = 50.0 + st = fresh_state() + order = [] + for _ in range(CL): # steps=CL -> k=1 per step + before = (st["canvas"][0] != mask_id).clone() + run_step(st, spiked, 0.0, leftmost=True, steps=CL) + new = ((st["canvas"][0] != mask_id) & ~before).nonzero() + order.extend(new.flatten().tolist()) + assert order == list(range(CL)), f"leftmost reveal order: {order}" + + +@pytest.mark.parametrize("max_steps,valid_length", [(4, 4), (1, 4), (4, 2)]) +def test_threshold_unmasking_and_commit(max_steps, valid_length, eager_sampler): + """Threshold decoding progresses, preserves revealed tokens, and clips padding.""" + import math + + from vllm.model_executor.models.nemotron_dllm import _compiled_masked_step + + device = "cpu" + width, vocab, mask = 4, 8, 7 + slots = torch.tensor([1], device=device) + indices = torch.tensor([0], device=device) + lengths = torch.tensor([valid_length], device=device) + temps = torch.zeros(1, device=device) + canvas = torch.full((2, width), mask, device=device, dtype=torch.int64) + step = torch.zeros(2, device=device, dtype=torch.int32) + phase = torch.zeros(2, device=device, dtype=torch.bool) + lp = torch.zeros(2, width, device=device) + ranks = torch.zeros(2, width, device=device, dtype=torch.int32) + sampled = torch.zeros(1, width, device=device, dtype=torch.int32) + count = torch.zeros(1, device=device, dtype=torch.int32) + draft = torch.zeros_like(canvas) + logits = torch.zeros(width, vocab, device=device) + logits[:, mask] = -100 + logits[0, 2] = 10 # Only the first position clears 0.9. + + def advance(): + _compiled_masked_step( + logits, + slots, + indices, + slots, + lengths, + temps, + canvas, + step, + phase, + lp, + ranks, + sampled, + count, + draft, + max_denoising_steps=max_steps, + mask_token_id=mask, + CL=width, + all_greedy=True, + capture_logprobs=True, + leftmost=False, + threshold_mode=True, + log_threshold=math.log(0.9), + ) + + advance() + assert count.item() == 0 + assert canvas[1, 0].item() == 2 + if max_steps > 1: + assert (canvas[1, 1:] == mask).all() + first_lp = lp[1, 0].clone() + logits[0, 3] = 20 # An already revealed position must remain unchanged. + for _ in range(valid_length - 1 if max_steps > 1 else 0): + advance() + assert phase[1].item() + assert (canvas[1, :valid_length] != mask).all() + assert (canvas[1, valid_length:] == mask).all() + assert canvas[1, 0].item() == 2 + assert lp[1, 0] == first_lp + completed = canvas[1].clone() + advance() + assert count.item() == valid_length + assert torch.equal(sampled[0, :valid_length], completed[:valid_length]) + assert (canvas[1] == mask).all() + assert not phase[1].item() + assert step[1].item() == 0 + assert (canvas[0] == mask).all() # Unscheduled request is untouched. + + +# --------------------------------------------------------------------------- +# HF reference (ground truth). transformers>=5 dropped two symbols the model's +# custom code imports; shim them to the shipped behavior. +# --------------------------------------------------------------------------- +def _install_tf_compat(): + import transformers.masking_utils as mu + import transformers.utils.generic as g + + if not hasattr(g, "check_model_inputs"): + + def check_model_inputs(func=None, **_): + return func if func is not None else (lambda f: f) + + g.check_model_inputs = check_model_inputs + if not hasattr(mu, "sdpa_mask_older_torch"): + mu.sdpa_mask_older_torch = mu.sdpa_mask + + +@pytest.fixture +def fp32_reference(monkeypatch): + """Isolate architecture parity from BF16 rounding and Triton TF32 dots.""" + precision = torch.get_float32_matmul_precision() + monkeypatch.setenv("VLLM_FLOAT32_MATMUL_PRECISION", "highest") + monkeypatch.setenv("TRITON_F32_DEFAULT", "ieee") + torch.set_float32_matmul_precision("highest") + yield + torch.set_float32_matmul_precision(precision) + + +def _load_tokenizer(): + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) + + +def _load_hf(dtype=torch.bfloat16, device="cuda"): + _install_tf_compat() + from transformers import AutoModel + + tok = _load_tokenizer() + model = ( + AutoModel.from_pretrained(MODEL_PATH, trust_remote_code=True, torch_dtype=dtype) + .to(device) + .eval() + ) + return tok, model + + +def _prompt_ids(tok, device="cuda"): + msgs = [{"role": "user", "content": "What is 2+2? Answer briefly."}] + enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt") + ids = enc["input_ids"] if not torch.is_tensor(enc) else enc + return ids.to(device) + + +def _vllm_manual_forward(obj, token_ids, causal_flag): + """Single forward of the loaded vLLM model over ``token_ids`` as one request + with per-request ``causal=causal_flag``, driving the GPU model runner + (KV-cache slots, attn metadata, forward context). Returns fp32 logits + ``[len(token_ids), vocab]`` on CPU. Runs under ``collective_rpc``. + """ + from vllm.forward_context import set_forward_context + from vllm.v1.worker.gpu.attn_utils import ( + build_attn_metadata, + build_slot_mappings_by_layer, + ) + + worker = getattr(obj, "worker", obj) + runner = worker.model_runner + device = runner.device + kv_cache_config = runner.kv_cache_config + assert len(kv_cache_config.kv_cache_groups) == 1 + block_size = kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size + + length = len(token_ids) + n_blocks = (length + block_size - 1) // block_size + # Block 0 is the null block; use fresh blocks from 1. + block_ids = torch.arange(1, n_blocks + 1, dtype=torch.int32, device=device) + block_table = block_ids.unsqueeze(0) + pos = torch.arange(length, dtype=torch.int64, device=device) + slot_mapping = ( + block_ids[pos // block_size].to(torch.int64) * block_size + pos % block_size + ) + slot_mappings = slot_mapping.unsqueeze(0) + + query_start_loc = torch.tensor([0, length], dtype=torch.int32, device=device) + query_start_loc_cpu = torch.tensor([0, length], dtype=torch.int32) + seq_lens = torch.tensor([length], dtype=torch.int32, device=device) + + causal = ( + True + if (isinstance(causal_flag, bool) and causal_flag) + else torch.tensor([bool(causal_flag)], dtype=torch.bool, device=device) + ) + + attn_metadata = build_attn_metadata( + attn_groups=runner.attn_groups, + num_reqs=1, + num_tokens=length, + query_start_loc_gpu=query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=length, + seq_lens=seq_lens, + max_seq_len=MAX_MODEL_LEN, + block_tables=[block_table], + slot_mappings=slot_mappings, + kv_cache_config=kv_cache_config, + causal=causal, + ) + slot_mappings_by_layer = build_slot_mappings_by_layer( + slot_mappings, kv_cache_config + ) + + input_ids = torch.tensor(token_ids, dtype=torch.int32, device=device) + with torch.inference_mode(): + with set_forward_context( + attn_metadata, + runner.vllm_config, + num_tokens=length, + slot_mapping=slot_mappings_by_layer, + ): + hidden = runner.model(input_ids=input_ids, positions=pos) + logits = runner.model.compute_logits(hidden) + return logits.float().cpu() + + +# --------------------------------------------------------------------------- +# GPU + weights: logits parity vs HF over a masked canvas (bidirectional) +# --------------------------------------------------------------------------- +@requires_gpu +@requires_weights +def test_logits_parity(monkeypatch, fp32_reference): + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + from vllm import LLM + from vllm.distributed import cleanup_dist_env_and_memory + + tok, hf = _load_hf(dtype=torch.float32) + prompt = _prompt_ids(tok) + canvas = torch.full( + (1, CANVAS_LENGTH), MASK_TOKEN_ID, dtype=prompt.dtype, device=prompt.device + ) + seq = torch.cat([prompt, canvas], dim=1) + token_ids = seq[0].tolist() + + with torch.no_grad(): + hf_logits = hf(seq).logits[0].float().cpu() # fully bidirectional + del hf + torch.accelerator.empty_cache() + + llm = LLM( + model=MODEL_PATH, + dtype="float32", + max_num_batched_tokens=128, + max_num_seqs=4, + attention_config={"backend": "TRITON_ATTN"}, + trust_remote_code=True, + enforce_eager=True, + max_model_len=MAX_MODEL_LEN, + gpu_memory_utilization=0.6, + ) + # causal_flag=False -> bidirectional canvas (matches HF single forward). + vllm_logits = llm.llm_engine.collective_rpc( + partial(_vllm_manual_forward, token_ids=token_ids, causal_flag=False) + )[0] + + llm.llm_engine.engine_core.shutdown() + del llm + cleanup_dist_env_and_memory() + + canvas_slice = slice(prompt.shape[1], seq.shape[1]) + hf_c = hf_logits[canvas_slice] + vl_c = vllm_logits[canvas_slice] + + argmax_match = (hf_c.argmax(-1) == vl_c.argmax(-1)).float().mean().item() + max_diff = (hf_c - vl_c).abs().max().item() + + assert argmax_match == 1.0, f"argmax mismatch on canvas: {argmax_match:.4f}" + assert max_diff < 1e-3, f"logit max|diff| too large: {max_diff:.6f}" + + +# --------------------------------------------------------------------------- +# GPU + weights: greedy end-to-end generation +# --------------------------------------------------------------------------- +@requires_gpu +@requires_weights +def test_greedy_generation(monkeypatch): + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + from vllm import LLM, SamplingParams + from vllm.distributed import cleanup_dist_env_and_memory + from vllm.inputs import TokensPrompt + + tok = _load_tokenizer() + prompt_ids = _prompt_ids(tok, device="cpu")[0].tolist() + + llm = LLM( + model=MODEL_PATH, + attention_config={"backend": "TRITON_ATTN"}, + trust_remote_code=True, + enforce_eager=True, + max_model_len=MAX_MODEL_LEN, + gpu_memory_utilization=0.6, + ) + out = llm.generate( + TokensPrompt(prompt_token_ids=prompt_ids), + SamplingParams(temperature=0.0, max_tokens=CANVAS_LENGTH), + ) + del llm + cleanup_dist_env_and_memory() + text = out[0].outputs[0].text + assert "4" in text, f"unexpected generation: {text!r}" + + +# --------------------------------------------------------------------------- +# GPU + weights: engine-wide stochastic sampling (RL rollouts) +# --------------------------------------------------------------------------- +@requires_gpu +@requires_weights +def test_stochastic_rollouts(monkeypatch): + """Mixed request temperatures override defaults and return reveal logprobs.""" + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + from vllm import LLM, SamplingParams + from vllm.distributed import cleanup_dist_env_and_memory + from vllm.inputs import TokensPrompt + + tok = _load_tokenizer() + msgs = [{"role": "user", "content": "Tell me a short story about a robot."}] + enc = tok.apply_chat_template(msgs, add_generation_prompt=True) + prompt_ids = enc["input_ids"] if not isinstance(enc, list) else enc + + llm = LLM( + model=MODEL_PATH, + attention_config={"backend": "TRITON_ATTN"}, + trust_remote_code=True, + enforce_eager=True, + max_model_len=MAX_MODEL_LEN, + gpu_memory_utilization=0.6, + diffusion_config={"model_options": {"temperature": 0.0}}, + ) + assert llm.get_default_sampling_params().temperature == 0.0 + temperatures = [0.0, 0.7, 1.0, 1.5] + outs = llm.generate( + [TokensPrompt(prompt_token_ids=list(prompt_ids))] * len(temperatures), + [ + SamplingParams( + temperature=t, top_p=0.9, max_tokens=2 * CANVAS_LENGTH, logprobs=0 + ) + for t in temperatures + ], + ) + del llm + cleanup_dist_env_and_memory() + texts = {o.outputs[0].text for o in outs} + assert len(texts) >= 2, f"Mixed-temperature rollouts were identical: {texts}" + for o in outs: + c = o.outputs[0] + assert c.logprobs is not None and len(c.logprobs) == len(c.token_ids) + for tid, entry in zip(c.token_ids, c.logprobs): + assert entry is not None + lp = entry[tid] + assert lp.rank is not None + assert lp.logprob <= 1e-6 and lp.rank >= 1 + + +def _assert_ar_runner(obj): + from vllm.v1.worker.gpu.model_states.default import DefaultModelState + from vllm.v1.worker.gpu.sample.sampler import Sampler + + runner = getattr(obj, "worker", obj).model_runner + assert type(runner.model_state) is DefaultModelState + assert type(runner.sampler) is Sampler + + +@requires_gpu +@requires_weights +@pytest.mark.parametrize( + "hf_overrides", + [ + {"ar_mode": True}, + {"architectures": ["NemotronLabsDiffusionForCausalLM"]}, + ], + ids=["ar_mode", "causal_architecture"], +) +def test_ar_logits_and_generation(monkeypatch, hf_overrides, fp32_reference): + """Both AR selectors use causal HF logits and ordinary cached generation.""" + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + from vllm import LLM, SamplingParams + from vllm.distributed import cleanup_dist_env_and_memory + from vllm.inputs import TokensPrompt + + tok, hf = _load_hf(dtype=torch.float32) + prompt = _prompt_ids(tok) + with torch.inference_mode(): + # Native AR generation disables diffusion_lm on the HF attention + # layers; otherwise they discard even an explicitly causal mask. + hf_output, _ = hf.ar_generate(prompt, max_new_tokens=16, temperature=0.0) + hf_logits = hf(prompt, use_causal_mask=True).logits[0].float().cpu() + hf_tokens = hf_output[0, prompt.shape[1] :].tolist() + del hf + torch.accelerator.empty_cache() + + llm = LLM( + model=MODEL_PATH, + dtype="float32", + max_num_batched_tokens=128, + max_num_seqs=4, + hf_overrides=hf_overrides, + attention_config={"backend": "TRITON_ATTN"}, + trust_remote_code=True, + enforce_eager=True, + max_model_len=MAX_MODEL_LEN, + gpu_memory_utilization=0.6, + ) + try: + config = llm.llm_engine.vllm_config + assert not config.model_config.is_diffusion + assert config.diffusion_config is None + assert config.speculative_config is None + llm.llm_engine.collective_rpc(_assert_ar_runner) + prompt_ids = prompt[0].tolist() + vllm_logits = llm.llm_engine.collective_rpc( + partial(_vllm_manual_forward, token_ids=prompt_ids, causal_flag=True) + )[0] + # Compare every prompt position, so a bidirectional prefill cannot pass + # by agreeing only on the last position. + torch.testing.assert_close(vllm_logits, hf_logits, atol=1e-3, rtol=0.0) + assert vllm_logits[-1].argmax() == hf_logits[-1].argmax() + + greedy = llm.generate( + TokensPrompt(prompt_token_ids=prompt_ids), + SamplingParams(temperature=0.0, max_tokens=16, logprobs=0), + )[0].outputs[0] + assert list(greedy.token_ids) == hf_tokens + assert greedy.logprobs is not None + # Exercise fractional-temperature sampling and logprobs as well. + sampled = llm.generate( + [TokensPrompt(prompt_token_ids=prompt_ids)] * 2, + SamplingParams( + temperature=0.7, top_p=0.9, top_k=20, max_tokens=8, logprobs=0 + ), + ) + for output in sampled: + completion = output.outputs[0] + assert 0 < len(completion.token_ids) <= 8 + assert completion.logprobs is not None + assert len(completion.logprobs) == len(completion.token_ids) + for token_id, entry in zip(completion.token_ids, completion.logprobs): + assert entry is not None + assert entry[token_id].logprob <= 1e-6 + finally: + llm.llm_engine.engine_core.shutdown() + del llm + torch._dynamo.reset() + cleanup_dist_env_and_memory() + + +def test_temperature_precedes_top_p(monkeypatch): + from vllm.model_executor.models import nemotron_dllm + from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch + + monkeypatch.setattr(nemotron_dllm, "apply_top_k_top_p", apply_top_k_top_p_pytorch) + logits = torch.tensor([[0.6, 0.3, 0.1, 10.0]]).log().repeat(3, 1) + filtered = nemotron_dllm._filter_sampling_logits( + logits, torch.tensor([0.0, 0.5, 2.0]), None, torch.full((3,), 0.7), 3 + ) + # At T=0.5 the leading token alone exceeds top_p; at T=1 or T=2 it + # takes two. The mask must not consume the nucleus probability budget. + expected = torch.tensor( + [ + [True, True, False, False], + [True, False, False, False], + [True, True, False, False], + ] + ) + assert torch.equal(filtered.isfinite(), expected) + torch.testing.assert_close(filtered[expected], logits[expected]) + + +def test_mixed_temperature_distribution_and_logprobs(eager_sampler): + from vllm.model_executor.models.nemotron_dllm import _compiled_masked_step + + # Many independent positions give a low-variance check of the actual + # sampling law, not just the temperatures attached to returned logprobs. + rows, width, vocab, mask = 3, 4096, 3, 2 + temperatures = torch.tensor([0.0, 0.5, 2.0]) + logits = torch.tensor([0.0, 1.0, float("-inf")]).expand(rows, width, vocab) + slots = torch.arange(rows) + canvas = torch.full((rows, width), mask, dtype=torch.int64) + logprobs = torch.zeros(rows, width) + torch.manual_seed(17) + processed = _compiled_masked_step( + logits.reshape(-1, vocab), + slots, + slots, + slots, + torch.full((rows,), width), + temperatures, + canvas, + torch.zeros(rows, dtype=torch.int32), + torch.zeros(rows, dtype=torch.bool), + logprobs, + torch.zeros(rows, width, dtype=torch.int32), + torch.zeros(rows, width, dtype=torch.int32), + torch.zeros(rows, dtype=torch.int32), + torch.zeros_like(canvas), + max_denoising_steps=1, + mask_token_id=mask, + CL=width, + all_greedy=False, + capture_logprobs=True, + leftmost=False, + threshold_mode=True, + log_threshold=0.0, + ) + assert (canvas[0] == 1).all() + for row, temperature in enumerate(temperatures.tolist()): + distribution = torch.distributions.Categorical( + logits=torch.tensor([0.0, 1.0]) / (temperature or 1.0) + ) + torch.testing.assert_close(logprobs[row], distribution.log_prob(canvas[row])) + torch.testing.assert_close( + processed[row].log_softmax(-1).gather(-1, canvas[row, :, None]).squeeze(-1), + logprobs[row], + ) + if temperature > 0: + frequency = (canvas[row] == 1).float().mean() + assert abs(frequency - distribution.probs[1]) < 0.03 diff --git a/tests/models/registry.py b/tests/models/registry.py index dbb97991a247..1b95e405f3e5 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -452,6 +452,14 @@ def check_available_online( trust_remote_code=True, is_available_online=False, ), + "NemotronLabsDiffusionForCausalLM": _HfExamplesInfo( + "nvidia/Nemotron-Labs-Diffusion-3B", + trust_remote_code=True, + hf_overrides={"architectures": ["NemotronLabsDiffusionForCausalLM"]}, + ), + "NemotronLabsDiffusionModel": _HfExamplesInfo( + "nvidia/Nemotron-Labs-Diffusion-3B", + ), "NanbeigeForCausalLM": _HfExamplesInfo( "Nanbeige/Nanbeige4.2-3B", trust_remote_code=True, diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index f32a3d4c367f..a2281da5ef8b 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, field from types import SimpleNamespace import pytest @@ -15,6 +15,7 @@ @dataclass class MockModelConfig: is_diffusion: bool = False + hf_config: SimpleNamespace = field(default_factory=SimpleNamespace) max_logprobs: int = 20 logits_processors: list | None = None return_sampling_mask: bool = False @@ -23,6 +24,29 @@ def get_vocab_size(self) -> int: return 1024 +def _admit( + params: SamplingParams, + hf_config: SimpleNamespace | None = None, + async_scheduling: bool = False, +) -> None: + """Run the engine's admission checks for a diffusion model.""" + processor = SimpleNamespace( + model_config=MockModelConfig( + is_diffusion=True, hf_config=hf_config or SimpleNamespace() + ), + vllm_config=SimpleNamespace( + scheduler_config=SimpleNamespace(async_scheduling=async_scheduling) + ), + speculative_config=None, + structured_outputs_config=None, + diffusion_config=DiffusionConfig(canvas_length=8), + tokenizer=None, + validate_logits_processors_params=lambda params: None, + resolve_watermarking=lambda params: False, + ) + InputProcessor._validate_params(processor, params, ("generate",)) + + @pytest.mark.parametrize( "kwargs", [ @@ -39,16 +63,23 @@ def get_vocab_size(self) -> int: def test_diffusion_rejects_unsupported_params(kwargs: dict): params = SamplingParams(**kwargs) with pytest.raises(VLLMValidationError, match="not yet supported with diffusion"): - params.verify(MockModelConfig(is_diffusion=True), None, None, None) + _admit(params) + + +@pytest.mark.parametrize("temperature", [0.0, 0.7, 1.5]) +def test_diffusion_per_request_temperature_is_a_model_opt_in(temperature): + hf_config = SimpleNamespace(diffusion_per_request_temperature=True) + _admit(SamplingParams(temperature=temperature), hf_config) + with pytest.raises(VLLMValidationError, match="not yet supported with diffusion"): + _admit(SamplingParams(temperature=temperature, seed=42), hf_config) def test_diffusion_accepts_default_params(): - SamplingParams().verify(MockModelConfig(is_diffusion=True), None, None, None) + _admit(SamplingParams()) def test_diffusion_accepts_top_k_top_p(): - params = SamplingParams(top_p=0.9, top_k=10) - params.verify(MockModelConfig(is_diffusion=True), None, None, None) + _admit(SamplingParams(top_p=0.9, top_k=10)) def test_non_diffusion_models_unaffected(): @@ -102,20 +133,7 @@ def test_diffusion_extra_args_are_validated_without_a_served_canvas(): def test_diffusion_canvas_width_is_accepted_with_either_scheduler( async_scheduling, extra_args ): - processor = SimpleNamespace( - model_config=MockModelConfig(is_diffusion=True), - vllm_config=SimpleNamespace( - scheduler_config=SimpleNamespace(async_scheduling=async_scheduling) - ), - speculative_config=None, - structured_outputs_config=None, - diffusion_config=DiffusionConfig(canvas_length=8), - tokenizer=None, - validate_logits_processors_params=lambda params: None, - resolve_watermarking=lambda params: False, - ) - params = SamplingParams(extra_args=extra_args) - InputProcessor._validate_params(processor, params, ("generate",)) + _admit(SamplingParams(extra_args=extra_args), async_scheduling=async_scheduling) @pytest.mark.parametrize( diff --git a/vllm/config/diffusion.py b/vllm/config/diffusion.py index 6f59c40a8366..63b6d0be8308 100644 --- a/vllm/config/diffusion.py +++ b/vllm/config/diffusion.py @@ -2,11 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Configuration for discrete diffusion (dLLM) models.""" +from typing import Any + from pydantic import Field from vllm.config.utils import config +@config +class DiffusionModelOptions: + """Base class for model-specific `DiffusionConfig.model_options`.""" + + @config class DiffusionConfig: """Configuration for discrete diffusion language models (dLLMs). @@ -24,3 +31,8 @@ class DiffusionConfig: max_denoising_steps: int | None = None """Maximum number of denoising iterations per canvas block. If not set, read from the model's generation_config.json.""" + + model_options: dict[str, Any] | DiffusionModelOptions = Field(default_factory=dict) + """Model-specific options. The accepted keys depend on the model and are + validated once the model is known, e.g. `NemotronDiffusionOptions` or + `DiffusionGemmaDiffusionOptions` in `vllm.model_executor.models.config`.""" diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index f796bcf6156a..d062c282ae03 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal, TypeVar +from pydantic import Field + +from vllm.config import utils as config_utils +from vllm.config.diffusion import DiffusionConfig, DiffusionModelOptions from vllm.logger import init_logger from vllm.utils.math_utils import round_up @@ -70,6 +74,27 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: ) +_OptionsT = TypeVar("_OptionsT", bound=DiffusionModelOptions) + + +def _parse_diffusion_model_options( + vllm_config: "VllmConfig", options_cls: type[_OptionsT] +) -> _OptionsT: + """Validate `--diffusion-config` `model_options` against `options_cls`.""" + diffusion_config = vllm_config.diffusion_config + assert diffusion_config is not None + options = diffusion_config.model_options + if isinstance(options, dict): + options = options_cls(**options) + if not isinstance(options, options_cls): + raise TypeError( + f"Expected {options_cls.__name__} diffusion model_options, " + f"got {type(options).__name__}." + ) + diffusion_config.model_options = options + return options + + class DeepseekV32ForCausalLM(VerifyAndUpdateConfig): @classmethod def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: @@ -386,6 +411,28 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: model_config.max_model_len = 8192 +@config_utils.config +class DiffusionGemmaDiffusionOptions(DiffusionModelOptions): + """DiffusionGemma `DiffusionConfig.model_options`. Unset options are read + from the model's generation_config.json.""" + + t_min: float | None = None + """Temperature at the last denoising step.""" + + t_max: float | None = None + """Temperature at the first denoising step.""" + + entropy_bound: float | None = Field(default=None, gt=0) + """Entropy budget for the EntropyBound sampler.""" + + confidence_threshold: float | None = None + """Confidence above which a position is committed.""" + + stability_threshold: int | None = Field(default=None, ge=0) + """Number of previous steps a prediction must match before it is + committed.""" + + class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig): @classmethod def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: @@ -421,14 +468,14 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: ) # Auto-create DiffusionConfig from HF config if not provided. + canvas_length = getattr( + vllm_config.model_config.hf_config, "canvas_length", 256 + ) if vllm_config.diffusion_config is None: - from vllm.config.diffusion import DiffusionConfig - - hf_config = vllm_config.model_config.hf_config - canvas_length = getattr(hf_config, "canvas_length", 256) - vllm_config.diffusion_config = DiffusionConfig( - canvas_length=canvas_length, - ) + vllm_config.diffusion_config = DiffusionConfig(canvas_length=canvas_length) + elif vllm_config.diffusion_config.canvas_length is None: + vllm_config.diffusion_config.canvas_length = canvas_length + _parse_diffusion_model_options(vllm_config, DiffusionGemmaDiffusionOptions) # The diffusion sampler materializes [num_seqs, canvas_length, vocab] # fp32 transients, so concurrency is memory-bound (>8 OOMs a single H200). @@ -457,6 +504,112 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: ) +@config_utils.config +class NemotronDiffusionOptions(DiffusionModelOptions): + """Nemotron Labs Diffusion `DiffusionConfig.model_options`.""" + + temperature: float | None = Field(default=None, ge=0) + """Default sampling temperature; explicit request temperatures take + precedence. 0 means greedy. If unset, the model's generation defaults + apply.""" + + selection_policy: ( + Literal["low_confidence", "leftmost", "confidence_threshold"] | None + ) = None + """Which masked positions to unmask each denoising step. + ``low_confidence``: the most confident predictions first (LLaDA), + top-k per step with k from the even transfer schedule. + ``leftmost``: strictly left-to-right. + ``confidence_threshold``: unmask every position whose chosen-token + probability exceeds ``confidence_threshold`` (always at least the single + most confident one, so each step makes progress). + If not set, read from the model's generation_config.json + (``diffusion_selection_policy``), defaulting to ``confidence_threshold``.""" + + confidence_threshold: float | None = Field(default=None, gt=0, lt=1) + """Probability threshold for ``selection_policy="confidence_threshold"``. + If not set, read from the model's generation_config.json + (``diffusion_confidence_threshold``), defaulting to 0.9.""" + + +class NemotronLabsDiffusionForBlockDiffusionConfig(VerifyAndUpdateConfig): + @classmethod + def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: + """Set up the diffusion config and defaults for the Nemotron + Ministral masked block-diffusion LM. + + Auto-creates DiffusionConfig from the HF config (canvas_length == + block_size) when the user didn't pass ``--diffusion-config``. + ``max_denoising_steps`` defaults to ``canvas_length`` at + sampler-build time when generation_config.json doesn't provide it. + """ + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + model_config = vllm_config.model_config + if not model_config.is_diffusion: + if vllm_config.diffusion_config is not None: + raise ValueError( + "Nemotron AR mode cannot be combined with --diffusion-config." + ) + if "max_new_tokens" not in model_config.override_generation_config: + model_config.override_generation_config["max_new_tokens"] = None + return + + attention_config = vllm_config.attention_config + if attention_config.backend == AttentionBackendEnum.FLASHINFER: + raise ValueError( + "FlashInfer does not support the Nemotron DLM's mixed " + "causal/bidirectional attention. Use --attention-backend " + "FLASH_ATTN or TRITON_ATTN instead." + ) + if attention_config.backend is None: + attention_config.backend = AttentionBackendEnum.TRITON_ATTN + logger.info("Nemotron Labs Diffusion defaults to TRITON_ATTN.") + elif attention_config.backend == AttentionBackendEnum.FLASH_ATTN: + if attention_config.flash_attn_version not in (None, 4): + raise ValueError( + "Nemotron Labs Diffusion requires FA4 for per-sequence " + "causal attention; use FA4 or TRITON_ATTN." + ) + attention_config.flash_attn_version = 4 + + # Auto-create DiffusionConfig from HF config if not provided; a + # user config passed just to set e.g. temperature backfills + # canvas_length from the HF config. + hf_config = vllm_config.model_config.hf_config + if vllm_config.diffusion_config is None: + vllm_config.diffusion_config = DiffusionConfig( + canvas_length=getattr(hf_config, "canvas_length", 32), + ) + elif vllm_config.diffusion_config.canvas_length is None: + vllm_config.diffusion_config.canvas_length = getattr( + hf_config, "canvas_length", 32 + ) + options = _parse_diffusion_model_options(vllm_config, NemotronDiffusionOptions) + + # Use the ordinary generation defaults; explicit request temperatures + # always take precedence over this compatibility setting. + if options.temperature is not None: + model_config.override_generation_config.setdefault( + "temperature", options.temperature + ) + + # The diffusion sampler materializes [num_seqs, canvas_length, vocab] + # fp32 transients, so cap concurrency when the user didn't pass + # --max-num-seqs (same policy as DiffusionGemma). + from vllm.config.scheduler import SchedulerConfig + + sc = vllm_config.scheduler_config + if sc is not None and sc.max_num_seqs >= SchedulerConfig.DEFAULT_MAX_NUM_SEQS: + sc.max_num_seqs = 8 + + # Remove any generation_config.json cap on max_new_tokens so each + # request controls its own output length via max_tokens. + model_config = vllm_config.model_config + if "max_new_tokens" not in model_config.override_generation_config: + model_config.override_generation_config["max_new_tokens"] = None + + class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -1219,6 +1372,8 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "NemotronHForCausalLM": NemotronHForCausalLMConfig, "NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig, "NemotronH_Nano_VL_V2": NemotronHNanoVLV2Config, + "NemotronLabsDiffusionForCausalLM": NemotronLabsDiffusionForBlockDiffusionConfig, + "NemotronLabsDiffusionModel": NemotronLabsDiffusionForBlockDiffusionConfig, "NomicBertModel": NomicBertModelConfig, "Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig, "Qwen2ForRewardModel": Qwen2ForRewardModelConfig, diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 7f4f9ffba99e..9711fa2b8579 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -37,6 +37,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, ) +from vllm.model_executor.models.config import DiffusionGemmaDiffusionOptions from vllm.model_executor.models.diffusion_gemma_sampler import ( sample_row_stats, sample_row_stats_reference, @@ -862,13 +863,25 @@ def __init__( 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 + assert diffusion_config is not None + options = diffusion_config.model_options + assert isinstance(options, DiffusionGemmaDiffusionOptions) + canvas_length = diffusion_config.canvas_length 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) + overrides = { + "t_min": options.t_min, + "t_max": options.t_max, + "confidence_threshold": options.confidence_threshold, + "stability_threshold": options.stability_threshold, + } + self.gen_config = { + **self.model_config.try_get_generation_config(), + **{k: v for k, v in overrides.items() if v is not None}, + } + max_denoising_steps = diffusion_config.max_denoising_steps or ( + self.gen_config.get("max_denoising_steps", 48) + ) self.diffusion_states = DiffusionGemmaRequestStates( max_num_reqs=self.max_num_reqs, canvas_length=canvas_length, @@ -910,11 +923,14 @@ def get_supported_generation_tasks(self): def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: diffusion_config = self.vllm_config.diffusion_config + assert diffusion_config is not None + options = diffusion_config.model_options + assert isinstance(options, DiffusionGemmaDiffusionOptions) 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") + entropy_bound = options.entropy_bound or 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})" diff --git a/vllm/model_executor/models/nemotron_dllm.py b/vllm/model_executor/models/nemotron_dllm.py new file mode 100644 index 000000000000..8488df9a5250 --- /dev/null +++ b/vllm/model_executor/models/nemotron_dllm.py @@ -0,0 +1,948 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Nemotron Labs Diffusion: masked block denoising or causal AR generation.""" + +from __future__ import annotations + +import math +from collections.abc import Iterable +from types import SimpleNamespace +from typing import Any + +import numpy as np +import torch +from torch import nn +from transformers import PreTrainedConfig + +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.models.config import NemotronDiffusionOptions +from vllm.model_executor.models.llama import ( + LlamaAttention, + LlamaDecoderLayer, + LlamaModel, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + maybe_prefix, +) +from vllm.sequence import IntermediateTensors +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.outputs import LogprobsTensors +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p +from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor +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_scores +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 SupportsPP, SupportsQuant + +logger = init_logger(__name__) + + +def _hf_yarn_get_mscale(scale: float, mscale: float = 1.0) -> float: + """``get_mscale`` from HF ``_compute_yarn_parameters``.""" + if scale <= 1.0: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class NemotronLabsDiffusionAttention(LlamaAttention): + """Llama attention + YaRN RoPE (HF attention factor) + per-token + llama-4 Q scaling applied after RoPE.""" + + def __init__(self, *, config: PreTrainedConfig, **kwargs: Any) -> None: + super().__init__(config=config, **kwargs) + rope_parameters = getattr(config, "rope_parameters", None) or {} + self.llama4_scaling_beta: float | None = rope_parameters.get( + "llama_4_scaling_beta" + ) + self.llama4_original_max_pos: int = rope_parameters.get( + "original_max_position_embeddings", + getattr(config, "max_position_embeddings", 16384), + ) + + def _init_rotary_emb( + self, + config: PreTrainedConfig, + quant_config: QuantizationConfig | None, + ) -> None: + rope_parameters = dict(getattr(config, "rope_parameters", None) or {}) + if rope_parameters.get("rope_type") == "yarn": + # The HF reference scales cos/sin by + # get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dim) + # (== 1.0 for this model), while vLLM's YaRN would apply + # 0.1*ln(factor)+1. Compute the HF factor and pass it via + # attn_factor with vLLM's own yarn scaling disabled. + if "attn_factor" not in rope_parameters: + factor = rope_parameters["factor"] + mscale = rope_parameters.get("mscale") + mscale_all_dim = rope_parameters.get("mscale_all_dim") + if mscale and mscale_all_dim: + attn_factor = _hf_yarn_get_mscale( + factor, mscale + ) / _hf_yarn_get_mscale(factor, mscale_all_dim) + else: + attn_factor = _hf_yarn_get_mscale(factor) + rope_parameters["attn_factor"] = float(attn_factor) + rope_parameters["apply_yarn_scaling"] = False + self.rotary_emb = get_rope( + self.head_dim, + max_position=self.max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=True, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + if self.llama4_scaling_beta is not None: + # scale = 1 + beta * log(1 + floor(pos / original_max_pos)), + # applied AFTER RoPE (matches HF `_get_llama_4_attn_scale`). + scale = 1.0 + self.llama4_scaling_beta * torch.log( + 1.0 + torch.floor(positions.float() / self.llama4_original_max_pos) + ) + q = q * scale.unsqueeze(-1).to(q.dtype) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + +class NemotronLabsDiffusionDecoderLayer(LlamaDecoderLayer): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + config: PreTrainedConfig | None = None, + ) -> None: + super().__init__( + vllm_config, + prefix=prefix, + config=config, + attn_layer_type=NemotronLabsDiffusionAttention, + ) + + +class NemotronLabsDiffusionForBlockDiffusion(nn.Module, SupportsPP, SupportsQuant): + """Ministral masked block-diffusion LM for vLLM. + + The backbone is a stock ``LlamaModel`` with Ministral parameters; the + only architectural deltas live in the attention layer (YaRN factor and + llama-4 Q scale). Mixed causal/bidirectional attention is driven by + ``NemotronLabsDiffusionModelState.prepare_attn`` via a per-request causal tensor. + """ + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "encoder.": "model.", + } + ) + + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + def get_model_state_cls(self): + if self.ar_mode: + from vllm.v1.worker.gpu.model_states.default import DefaultModelState + + return DefaultModelState + return NemotronLabsDiffusionModelState + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.ar_mode = not vllm_config.model_config.is_diffusion + + self.model = LlamaModel( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + layer_type=NemotronLabsDiffusionDecoderLayer, + ) + # The LM head is named diffusion_head in the checkpoint and is NOT + # tied to the embeddings. + self.diffusion_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "diffusion_head"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.diffusion_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + +@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) + + +def _filter_sampling_logits( + logits: torch.Tensor, + temperatures: torch.Tensor, + top_k: torch.Tensor | None, + top_p: torch.Tensor | None, + mask_token_id: int, +) -> torch.Tensor: + """Select sampling candidates after temperature scaling, excluding masks.""" + logits = logits.float().clone() + logits[:, mask_token_id] = float("-inf") + temperatures = torch.where(temperatures > 0, temperatures, 1.0) + filtered = apply_top_k_top_p(logits / temperatures[:, None], top_k, top_p) + # Retain unscaled logits for the confidence-based reveal policy. + return logits.masked_fill_(filtered.isneginf(), float("-inf")) + + +@torch.compile(dynamic=True) +def _compiled_masked_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 len (<=CL) + req_temps: torch.Tensor, # [num_decode] fp32 per-request temperature (0=greedy) + # State tensors (modified in-place) + canvas: torch.Tensor, # [max_num_reqs, CL] + step_tensor: torch.Tensor, # [max_num_reqs] + is_encoder_phase: torch.Tensor, # [max_num_reqs] + unmask_logprobs: torch.Tensor, # [max_num_reqs, CL] fp32 + unmask_ranks: torch.Tensor, # [max_num_reqs, CL] int32 + # 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: int, + mask_token_id: int, + CL: int, + all_greedy: bool, + capture_logprobs: bool, + leftmost: bool, + threshold_mode: bool, + log_threshold: float, +) -> torch.Tensor: + """Update absorbing masked-token state and emit completed blocks. + + Confidence uses the selected token's probability before temperature scaling. + Capture token logprobs when positions are revealed, then refresh causal KV + on the following commit step. Return logits for top-k introspection. + """ + num_decode = decode_slots.shape[0] + device = decode_slots.device + + logits_3d = logits.reshape(num_decode, CL, -1).float().clone() + logits_3d[..., mask_token_id] = float("-inf") + + is_commit = is_encoder_phase[decode_slots] # [num_decode] + is_denoise = ~is_commit + + cur_canvas = canvas[decode_slots] # [num_decode, CL] + pos = torch.arange(CL, device=device) + # A canvas truncated near max_model_len is zero-padded up to CL by the + # caller; the padded tail is treated as non-masked so it is never + # transferred or committed. + valid = pos.unsqueeze(0) < valid_canvas_len.unsqueeze(1) # [num_decode, CL] + mask_index = (cur_canvas == mask_token_id) & valid + + # ---- x0 + confidence (prob of chosen token) at masked positions ---- + if not all_greedy: + # Gumbel-max: argmax(logits + T*g) ~ softmax(logits / T). + # Rows whose request asked for greedy (per-request temperature 0) + # get zero noise and reduce to argmax. + u = torch.rand_like(logits_3d).clamp_min(1e-20) + gumbel = -torch.log(-torch.log(u)) + x0 = (logits_3d + req_temps[:, None, None] * gumbel).argmax(dim=-1) + else: + x0 = logits_3d.argmax(dim=-1) # [num_decode, CL] + chosen_logits = logits_3d.gather(-1, x0.unsqueeze(-1)).squeeze(-1) + x0_logprob = chosen_logits - logits_3d.logsumexp(dim=-1) + neg_inf = torch.full_like(x0_logprob, float("-inf")) + confidence = torch.where(mask_index, x0_logprob, neg_inf) + + # ---- Even transfer schedule: k = ceil(#masks / remaining steps) ---- + num_masks = mask_index.sum(dim=-1) # [num_decode] int64 + cur_step = step_tensor[decode_slots].to(torch.int64) + remaining_steps = (max_denoising_steps - cur_step).clamp(min=1) + k = -(-num_masks // remaining_steps) # ceil div + + if leftmost: + # Leftmost-reveal: unmask the first k masked positions (strict + # left-to-right order instead of confidence order). + csum = torch.cumsum(mask_index.to(torch.int64), dim=-1) + transfer = mask_index & (csum <= k.unsqueeze(1)) + elif threshold_mode: + # FastDiffuser thresholding: unmask every masked position whose + # chosen-token logprob clears the threshold, and always the single + # most confident one (guaranteed progress; also covers the case + # where nothing clears the bar). `confidence` is -inf outside + # mask_index, so argmax lands on a masked position whenever one + # exists and both terms are automatically gated by mask_index. + above = confidence >= log_threshold + top1 = confidence.argmax(dim=-1, keepdim=True) # [num_decode, 1] + force = torch.zeros_like(mask_index).scatter_(1, top1, True) + transfer = mask_index & (above | force | (remaining_steps == 1).unsqueeze(1)) + else: + # Top-k per row with per-row k: rank positions by confidence (desc). + order = torch.argsort(confidence, dim=-1, descending=True, stable=True) + ranks = torch.empty_like(order) + ranks.scatter_(1, order, pos.expand(num_decode, CL)) + transfer = mask_index & (ranks < k.unsqueeze(1)) + + # ---- At-unmask logprob/rank capture (sampling distribution) ---- + sample_logits = logits_3d + if capture_logprobs: + # Greedy requests report unscaled model logprobs, as in AR sampling. + row_temp = torch.where(req_temps > 0, req_temps, 1.0) + sample_logits = logits_3d / row_temp[:, None, None] + chosen_lp = sample_logits.gather(-1, x0.unsqueeze(-1)).squeeze( + -1 + ) - sample_logits.logsumexp(dim=-1) + chosen_rank = ( + (logits_3d >= chosen_logits.unsqueeze(-1)).sum(dim=-1).to(torch.int32) + ) + write = transfer & is_denoise.unsqueeze(1) + unmask_logprobs[decode_slots] = torch.where( + write, chosen_lp, unmask_logprobs[decode_slots] + ) + unmask_ranks[decode_slots] = torch.where( + write, chosen_rank, unmask_ranks[decode_slots] + ) + + # ---- Absorbing-state update (denoise rows only) ---- + new_canvas = torch.where(transfer & is_denoise.unsqueeze(1), x0, cur_canvas) + + # Commit rows: reset the canvas to masks for the next block. + mask_canvas = torch.full_like(cur_canvas, mask_token_id) + canvas[decode_slots] = torch.where(is_commit.unsqueeze(1), mask_canvas, new_canvas) + + # Step counter: +1 for denoise, reset to 0 for commit. + new_step = torch.where( + is_denoise, + (cur_step + 1).to(step_tensor.dtype), + step_tensor.new_zeros(num_decode), + ) + step_tensor[decode_slots] = new_step + + # Sampled output: the commit step emits the canvas that was fed to the + # model this step; denoise steps emit nothing. A canvas truncated by + # max_model_len can converge with unresolved masks in its tail (only the + # scheduled prefix is denoised); never emit mask tokens — commit stops at + # the first remaining mask (for leftmost reveal this is exactly the clean + # prefix). + mask_in_canvas = (cur_canvas == mask_token_id) & valid + first_mask = torch.where( + mask_in_canvas.any(dim=-1), + mask_in_canvas.to(torch.int64).argmax(dim=-1), + valid_canvas_len, + ) + emit_len = torch.minimum(valid_canvas_len, first_mask) + sampled[decode_idx] = cur_canvas.to(sampled.dtype) * is_commit.unsqueeze(1).to( + sampled.dtype + ) + num_sampled[decode_idx] = is_commit.to(num_sampled.dtype) * emit_len.to( + num_sampled.dtype + ) + + # ---- Convergence: no masks remain, or step budget exhausted ---- + masks_left = ((new_canvas == mask_token_id) & valid).sum(dim=-1) + converged = (masks_left == 0) | (new_step.to(torch.int64) >= max_denoising_steps) + # Commit done -> denoise next (False); denoise converged -> commit (True). + is_encoder_phase[decode_slots] = torch.where( + is_commit, is_commit.new_zeros(num_decode), converged + ) + + # Feed the canvas back to the scheduler as next step's draft tokens. + draft_tokens[all_slots, :CL] = canvas[all_slots] + + return sample_logits + + +class NemotronLabsDiffusionRequestStates: + """Pre-allocated GPU tensors for per-request masked-diffusion state.""" + + def __init__( + self, + max_num_reqs: int, + canvas_length: int, + mask_token_id: int, + max_denoising_steps: int, + device: torch.device, + ): + self.max_num_reqs = max_num_reqs + self.canvas_length = canvas_length + self.mask_token_id = mask_token_id + self.max_denoising_steps = max_denoising_steps + self.device = device + + self.is_encoder_phase = torch.zeros( + max_num_reqs, dtype=torch.bool, device=device + ) + # Canvas tokens [max_num_reqs, canvas_length], mask_token_id when + # not yet unmasked. + self.canvas = torch.full( + (max_num_reqs, canvas_length), + mask_token_id, + dtype=torch.int64, + device=device, + ) + # Denoising step counter (0..max_denoising_steps). + self.step = torch.zeros(max_num_reqs, dtype=torch.int32, device=device) + # Per-position logprob/rank of the chosen token, captured at the + # denoise step each position was unmasked; emitted at commit. + self.unmask_logprobs = torch.zeros( + max_num_reqs, canvas_length, dtype=torch.float32, device=device + ) + self.unmask_ranks = torch.zeros( + max_num_reqs, canvas_length, dtype=torch.int32, device=device + ) + + def init_canvas(self, slot_indices) -> None: + """Reset the canvas to all-mask for the given slots.""" + self.canvas[slot_indices] = self.mask_token_id + + def add_request(self, slot_idx: int) -> None: + self.is_encoder_phase[slot_idx] = True + self.canvas[slot_idx] = self.mask_token_id + self.step[slot_idx] = 0 + self.unmask_logprobs[slot_idx] = 0.0 + self.unmask_ranks[slot_idx] = 0 + + def remove_request(self, slot_idx: int) -> None: + self.is_encoder_phase[slot_idx] = False + + +class NemotronLabsDiffusionModelState(ModelState): + """ModelState for the Ministral masked block-diffusion LM. + + Single backbone in two modes: + - encoder mode (prefill/commit): causal attention, writes KV + - denoise mode: bidirectional attention over the canvas + """ + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: Any, + device: torch.device, + ) -> None: + super().__init__(vllm_config, model, encoder_cache, device) + + hf_config = self.model_config.hf_config + diffusion_config = vllm_config.diffusion_config + assert diffusion_config is not None + options = diffusion_config.model_options + assert isinstance(options, NemotronDiffusionOptions) + canvas_length = diffusion_config.canvas_length + gen_config = self.model_config.try_get_generation_config() + # Default to one unmasked token per step (HF steps == gen_length). + max_denoising_steps = ( + diffusion_config.max_denoising_steps + or gen_config.get("max_denoising_steps") + or canvas_length + ) + selection = ( + options.selection_policy + or gen_config.get("diffusion_selection_policy") + or "confidence_threshold" + ) + if selection not in ("low_confidence", "leftmost", "confidence_threshold"): + raise ValueError(f"Unknown diffusion selection_policy: {selection!r}") + self.selection_policy = selection + # Probability threshold for selection_policy="confidence_threshold" + # (SGLang FastDiffuser thresholding). None-checks rather than `or` + # so an explicit small value is honored. + threshold = options.confidence_threshold + if threshold is None: + threshold = gen_config.get("diffusion_confidence_threshold") + if threshold is None: + threshold = 0.9 + self.confidence_threshold = float(threshold) + self.diffusion_states = NemotronLabsDiffusionRequestStates( + max_num_reqs=self.max_num_reqs, + canvas_length=canvas_length, + mask_token_id=getattr(hf_config, "mask_token_id", 100), + max_denoising_steps=max_denoising_steps, + device=device, + ) + 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 + ) + + def get_supported_generation_tasks(self): + return ("generate",) + + def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: + return MaskedDiffusionSampler( + sampler=sampler, + vocab_size=self.model_config.get_vocab_size(), + diffusion_states=self.diffusion_states, + selection_policy=self.selection_policy, + confidence_threshold=self.confidence_threshold, + ), None + + 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) + + 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 prepare_inputs_embeds( + self, + scheduled_encoder_inputs: dict[str, list[int]], + input_batch: InputBatch, + req_states: RequestState, + ) -> torch.Tensor | None: + return None + + def prepare_inputs(self, input_batch, req_states) -> dict[str, Any]: + # Text-only model, no self-conditioning: use the default + # input_ids path. + return {} + + def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: + return {} + + def prepare_attn( + self, + input_batch, + cudagraph_mode, + block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + for_capture=False, + ubatch_idx: int = 0, + ) -> dict[str, Any]: + assert ubatch_idx == 0, "DBO is not supported" + 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 (prefill/commit) = causal, + # denoise = bidirectional. Pass a 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 + # 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 MaskedDiffusionSampler: + """Masked diffusion with per-request sampling temperatures. + + The default threshold policy reveals all confident positions and guarantees + at least one reveal per step. ``logprobs=0`` returns each token's logprob at + reveal time; top-k introspection uses the converged denoising distribution. + """ + + def __init__( + self, + sampler: Any, + vocab_size: int, + diffusion_states: NemotronLabsDiffusionRequestStates, + selection_policy: str = "confidence_threshold", + confidence_threshold: float = 0.9, + ): + self.sampling_states = sampler.sampling_states + self.req_states = sampler.req_states + self.vocab_size = vocab_size + self.diffusion_states = diffusion_states + self.canvas_length = diffusion_states.canvas_length + self.mask_token_id = diffusion_states.mask_token_id + self.leftmost = selection_policy == "leftmost" + self.threshold_mode = selection_policy == "confidence_threshold" + # Compared against per-position chosen-token logprobs (natural log). + self.log_threshold = math.log(confidence_threshold) + + 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 top-k (logprobs>0) tensors computed on the + # converging denoise step; consumed on the subsequent commit step. + self._pending_logprobs: dict[int, LogprobsTensors] = {} + + def add_request(self, req_idx: int, sampling_params: Any) -> None: + if use_penalty(sampling_params): + logger.warning_once( + "Masked block-diffusion 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): + 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 all-mask 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 + ps_gpu = async_tensor_h2d( + ps.astype(np.int64), device=states.is_encoder_phase.device + ) + states.init_canvas(ps_gpu) + self.req_states.draft_tokens[ps_gpu, : self.canvas_length] = states.canvas[ + ps_gpu + ] + states.is_encoder_phase.index_fill_(0, ps_gpu, False) + + def _handle_prefill(self, input_batch: Any) -> 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, + logprobs_tensors: LogprobsTensors | None = None, + ) -> 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) + + # --- CPU/NumPy setup (outside compile): split decode vs prefill, + # init canvas for any new prefills, stage decode slot indices. --- + 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; == CL except for a canvas + # truncated near max_model_len. + valid_canvas_len_np = per_req_nlogits_np[per_req_nlogits_np > 0] + valid_canvas_len = async_tensor_h2d( + valid_canvas_len_np.astype(np.int64), device=device + ) + + if num_decode > 0: + token_slots = decode_slots.repeat_interleave( + valid_canvas_len, output_size=int(valid_canvas_len_np.sum()) + ) + top_k, top_p = self.sampling_states.get_top_k_top_p( + token_slots, decode_slots_np + ) + if top_k is not None or top_p is not None: + logits = _filter_sampling_logits( + logits, + self.sampling_states.temperature.gpu[token_slots], + top_k, + top_p, + self.mask_token_id, + ) + + # Pad any truncated canvas back to CL so the uniform-CL sampler math + # holds. Padded positions are zeroed (uniform logits) and are treated + # as non-masked (never transferred or committed). + 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 + valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1) + src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1) + logits = logits[src.reshape(-1)].masked_fill(~valid.reshape(-1, 1), 0) + + 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() + + max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np) + + req_temps = self.sampling_states.temperature.gpu[decode_slots] + + logits_3d = _compiled_masked_step( + logits[: num_decode * CL], + decode_slots, + decode_idx, + all_slots, + valid_canvas_len, + req_temps, + # State + states.canvas, + states.step, + states.is_encoder_phase, + states.unmask_logprobs, + states.unmask_ranks, + # Output + sampled, + num_sampled, + self.req_states.draft_tokens, + # Config + max_denoising_steps=int(states.max_denoising_steps), + mask_token_id=int(states.mask_token_id), + CL=CL, + all_greedy=bool( + (self.sampling_states.temperature.np[decode_slots_np] == 0).all() + ), + capture_logprobs=max_num_logprobs >= 0, + leftmost=self.leftmost, + threshold_mode=self.threshold_mode, + log_threshold=self.log_threshold, + ) + + # Top-k introspection (logprobs>0): stash the converged denoise + # step's top-k (is_encoder_phase flipped False->True); attached on + # the commit. The logprobs==0 path reads the at-unmask buffers + # instead and needs no stash. + if max_num_logprobs > 0 and num_decode > 0: + converged_mask = states.is_encoder_phase[decode_slots] + just_converged = converged_mask & ~is_committing + if just_converged.any(): + flat_logits = logits_3d.reshape(-1, logits_3d.shape[-1]) + final_canvas = states.canvas[decode_slots] + for local_idx in just_converged.nonzero(as_tuple=True)[0]: + li = int(local_idx.item()) + slot = int(decode_slots[local_idx].item()) + k_i = int(valid_canvas_len_np[li]) + row = li * CL + self._pending_logprobs[slot] = compute_topk_scores( + flat_logits[row : row + k_i], + max_num_logprobs, + final_canvas[local_idx][:k_i], + ) + + # Commit steps: emit logprobs for the committed block. The sampled + # token's logprob/rank come from the at-unmask buffers when only + # they are needed (logprobs==0); top-k rows come from the stash. + logprobs_tensors = None + emit = max_num_logprobs >= 0 and num_decode > 0 + committing_np = is_committing.cpu().numpy() if emit else None + if emit and committing_np is not None and committing_np.any(): + decode_row_of_req = { + int(decode_indices_np[li]): li for li in range(num_decode) + } + parts_ids, parts_lp, parts_ranks = [], [], [] + cu_gen: list[int] = [] + flat_offset = 0 + for i in range(num_reqs): + cu_gen.append(flat_offset) + decode_row = decode_row_of_req.get(i) + if decode_row is None or not committing_np[decode_row]: + continue + slot = int(slots_np[i]) + k_i = int(valid_canvas_len_np[decode_row]) + if max_num_logprobs > 0: + lp = self._pending_logprobs.pop(slot, None) + if lp is None: + continue + 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] + else: + parts_ids.append(sampled[i, :k_i].to(torch.int64).unsqueeze(-1)) + parts_lp.append(states.unmask_logprobs[slot, :k_i].unsqueeze(-1)) + parts_ranks.append(states.unmask_ranks[slot, :k_i]) + flat_offset += k_i + 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, + logprobs_tensors=logprobs_tensors, + ) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 23e67c9f506c..6db689e51100 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -179,6 +179,14 @@ "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"), "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), + "NemotronLabsDiffusionForCausalLM": ( + "nemotron_dllm", + "NemotronLabsDiffusionForBlockDiffusion", + ), + "NemotronLabsDiffusionModel": ( + "nemotron_dllm", + "NemotronLabsDiffusionForBlockDiffusion", + ), "MuseGlimmerForCausalLM": ("muse_glimmer", "MuseGlimmerForCausalLM"), "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"), "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"), diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 8331a1707b7d..702576368e2a 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -911,7 +911,6 @@ def verify( self._validate_stop_token_ids(model_config) self._validate_allowed_token_ids(model_config) self._validate_spec_decode(speculative_config) - self._validate_diffusion(model_config) self._validate_structured_outputs( model_config, structured_outputs_config, tokenizer ) @@ -1193,28 +1192,6 @@ def _validate_spec_decode( "are not yet supported with speculative decoding." ) - def _validate_diffusion(self, model_config: ModelConfig) -> None: - if not model_config.is_diffusion: - return - - # Diffusion models denoise a whole canvas per step with a fixed - # temperature schedule, so per-request sampling parameters are not - # supported. Penalties are ignored by the sampler with a warning. - if ( - self.temperature != 1.0 - or self.min_p > _SAMPLING_EPS - or self.seed is not None - or self.min_tokens > 0 - or self.logit_bias - or self.bad_words - or self.allowed_token_ids - ): - raise VLLMValidationError( - "The temperature, min_p, seed, min_tokens, logit_bias, " - "bad_words, and allowed_token_ids sampling parameters " - "are not yet supported with diffusion models." - ) - def _validate_structured_outputs( self, model_config: ModelConfig, diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index c67d468a13b3..2bad6ae2a5e4 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -144,6 +144,7 @@ def __getitem__(self, key): kimi_k3="KimiK3Config", RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct) RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct) + nemotron_labs_diffusion="NemotronLabsDiffusionConfig", midashenglm="MiDashengLMConfig", minimax_m3_vl="MiniMaxM3Config", minimax_m3_mtp="MiniMaxM3MTPConfig", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 68f82ead6b39..bb037c78939a 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -61,6 +61,9 @@ "MuseGlimmerAssistantConfig": "vllm.transformers_utils.configs.muse_glimmer", "KimiK3Config": "vllm.transformers_utils.configs.kimi_k3", "KimiK3VisionConfig": "vllm.transformers_utils.configs.kimi_k3", + "NemotronLabsDiffusionConfig": ( + "vllm.transformers_utils.configs.nemotron_labs_diffusion" + ), "OpenVLAConfig": "vllm.transformers_utils.configs.openvla", "OvisConfig": "vllm.transformers_utils.configs.ovis", "PixelShuffleSiglip2VisionConfig": "vllm.transformers_utils.configs.isaac", @@ -116,6 +119,7 @@ "MiniMaxM3Config", "MiniMaxM3MTPConfig", "MiniMaxM3TextConfig", + "NemotronLabsDiffusionConfig", "Moondream3Config", "Moondream3TextConfig", "Moondream3VisionConfig", diff --git a/vllm/transformers_utils/configs/nemotron_labs_diffusion.py b/vllm/transformers_utils/configs/nemotron_labs_diffusion.py new file mode 100644 index 000000000000..78ea7076d5ab --- /dev/null +++ b/vllm/transformers_utils/configs/nemotron_labs_diffusion.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Config for Nemotron Ministral-based masked block-diffusion LMs. + +The HF checkpoint ships a trust-remote-code ``NemotronLabsDiffusionConfig`` that does +not expose ``canvas_length``, which vLLM uses to detect diffusion models +(``ModelConfig.is_diffusion``). This local config mirrors the Ministral +backbone fields and derives ``canvas_length`` from ``block_size``. +""" + +from typing import Any + +from transformers import PreTrainedConfig + + +class NemotronLabsDiffusionConfig(PreTrainedConfig): + model_type = "nemotron_labs_diffusion" + + def __init__( + self, + vocab_size: int = 131072, + hidden_size: int = 3072, + intermediate_size: int = 9216, + num_hidden_layers: int = 26, + num_attention_heads: int = 32, + num_key_value_heads: int = 8, + head_dim: int = 128, + hidden_act: str = "silu", + max_position_embeddings: int = 262144, + initializer_range: float = 0.02, + rms_norm_eps: float = 1e-5, + use_cache: bool = False, + rope_theta: float = 1000000.0, + rope_parameters: dict[str, Any] | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + mlp_bias: bool = False, + sliding_window: int | None = None, + mask_token_id: int = 100, + block_size: int = 32, + **kwargs: Any, + ) -> None: + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.mlp_bias = mlp_bias + self.sliding_window = sliding_window + + # Masked block-diffusion fields. ``canvas_length`` (a property below) + # drives vLLM's diffusion detection and the per-step canvas size. + self.mask_token_id = mask_token_id + self.block_size = block_size + # Nemotron accepts per-request temperatures in diffusion mode, unlike + # diffusion models that use an engine-level temperature schedule. + self.diffusion_per_request_temperature = True + + # Normalize RoPE config: the checkpoint carries the same YaRN dict + # under both ``rope_parameters`` and (legacy) ``rope_scaling``. Pop + # the legacy key so PreTrainedConfig's rope_scaling property setter + # doesn't overwrite the normalized dict below. + rope_scaling = kwargs.pop("rope_scaling", None) + rope_parameters = dict(rope_parameters or rope_scaling or {}) + rope_parameters.setdefault("rope_theta", rope_theta) + rope_parameters.setdefault("rope_type", rope_parameters.pop("type", "default")) + self.rope_parameters = rope_parameters + self.rope_theta = rope_theta + + super().__init__( + pad_token_id=kwargs.pop("pad_token_id", None), + tie_word_embeddings=kwargs.pop("tie_word_embeddings", False), + **kwargs, + ) + + @property + def canvas_length(self) -> int | None: + """Canvas width for masked diffusion, or ``None`` for the AR path. + + The same checkpoint can be served autoregressively, either with + ``hf_overrides={"ar_mode": True}`` or through the causal-LM + architecture. ``ModelConfig.is_diffusion`` keys off this value. + """ + if getattr(self, "ar_mode", False): + return None + if "NemotronLabsDiffusionForCausalLM" in ( + getattr(self, "architectures", None) or () + ): + return None + return self.block_size + + @canvas_length.setter + def canvas_length(self, value: int | None) -> None: + if value is not None: + self.block_size = value diff --git a/vllm/utils/diffusion.py b/vllm/utils/diffusion.py index f3018c18e7bd..b8aae6a6ec77 100644 --- a/vllm/utils/diffusion.py +++ b/vllm/utils/diffusion.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING from vllm.exceptions import VLLMValidationError +from vllm.sampling_params import _SAMPLING_EPS if TYPE_CHECKING: from vllm.sampling_params import SamplingParams @@ -15,8 +16,30 @@ def validate_diffusion_sampling_params( *, canvas_length: int | None, vocab_size: int, + per_request_temperature: bool = False, ) -> None: - """Validate request options and normalize one-canvas reads at admission.""" + """Validate request options and normalize one-canvas reads at admission. + + Diffusion models denoise a whole canvas per step, so most per-request + sampling parameters are not supported. Models whose sampler applies the + request temperature set `per_request_temperature`. Penalties are ignored + by the sampler with a warning. + """ + if ( + (params.temperature != 1.0 and not per_request_temperature) + or params.min_p > _SAMPLING_EPS + or params.seed is not None + or params.min_tokens > 0 + or params.logit_bias + or params.bad_words + or params.allowed_token_ids + ): + raise VLLMValidationError( + "The temperature, min_p, seed, min_tokens, logit_bias, " + "bad_words, and allowed_token_ids sampling parameters " + "are not yet supported with diffusion models." + ) + extra = params.extra_args if not extra: return diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 17e8af5caeef..a5297cf92d51 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -187,6 +187,11 @@ def _validate_params( else None ), vocab_size=self.model_config.get_vocab_size(), + per_request_temperature=getattr( + self.model_config.hf_config, + "diffusion_per_request_temperature", + False, + ), ) if self.model_config.return_sampling_mask: