diff --git a/tests/models/language/generation/test_nemotron_dllm.py b/tests/models/language/generation/test_nemotron_dllm.py index 70d4c43c533e..e0bc878cf8a0 100644 --- a/tests/models/language/generation/test_nemotron_dllm.py +++ b/tests/models/language/generation/test_nemotron_dllm.py @@ -101,7 +101,10 @@ def test_arch_in_model_registry(): @pytest.mark.parametrize("flash_version", [None, 3]) -def test_attention_backend_supports_mixed_causality(flash_version): +@pytest.mark.parametrize("pipeline_parallel_size", [1, 2]) +def test_attention_backend_supports_mixed_causality( + flash_version, pipeline_parallel_size +): from types import SimpleNamespace from typing import cast @@ -128,9 +131,17 @@ def test_attention_backend_supports_mixed_causality(flash_version): ), diffusion_config=None, scheduler_config=None, + parallel_config=SimpleNamespace( + pipeline_parallel_size=pipeline_parallel_size + ), ), ) - if flash_version is not None: + if pipeline_parallel_size != 1: + with pytest.raises(ValueError, match="PP=1"): + NemotronLabsDiffusionForBlockDiffusionConfig.verify_and_update_config( + config + ) + elif flash_version is not None: with pytest.raises(ValueError, match="requires FA4"): NemotronLabsDiffusionForBlockDiffusionConfig.verify_and_update_config( config @@ -707,6 +718,72 @@ def test_stochastic_rollouts(monkeypatch): assert lp.logprob <= 1e-6 and lp.rank >= 1 +@requires_gpu +@requires_weights +@pytest.mark.parametrize("enforce_eager", [True, False]) +def test_diffusion_logprobs_independent_of_batch_requests(monkeypatch, enforce_eager): + """A neighbour's top-k request must not change sampled-token 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() + prompt = TokensPrompt(prompt_token_ids=_prompt_ids(tok, device="cpu")[0].tolist()) + budgets = (17, 41, 9) + llm = LLM( + model=MODEL_PATH, + attention_config={"backend": "TRITON_ATTN"}, + trust_remote_code=True, + enforce_eager=enforce_eager, + max_model_len=MAX_MODEL_LEN, + max_num_seqs=4, + gpu_memory_utilization=0.6, + diffusion_config={"model_options": {"selection_policy": "leftmost"}}, + ) + try: + assert llm.reset_prefix_cache() + baseline = llm.generate( + [prompt] * 3, + [ + SamplingParams(temperature=0, max_tokens=n, logprobs=0, ignore_eos=True) + for n in budgets + ], + use_tqdm=False, + ) + assert llm.reset_prefix_cache() + mixed = llm.generate( + [prompt] * 3, + [ + SamplingParams(temperature=0, max_tokens=n, logprobs=k, ignore_eos=True) + for n, k in zip(budgets, (None, 0, 2)) + ], + use_tqdm=False, + ) + assert mixed[0].outputs[0].logprobs is None + assert [len(out.outputs[0].token_ids) for out in mixed] == list(budgets) + for reference, result in zip(baseline[1:], mixed[1:]): + expected_completion = reference.outputs[0] + completion = result.outputs[0] + assert expected_completion.logprobs is not None + assert completion.token_ids == expected_completion.token_ids + assert completion.logprobs is not None + assert len(completion.logprobs) == len(completion.token_ids) + for token, expected, actual in zip( + expected_completion.token_ids, + expected_completion.logprobs, + completion.logprobs, + ): + assert expected is not None and token in expected + assert actual is not None and token in actual + assert actual[token].logprob == pytest.approx( + expected[token].logprob, abs=0.02 + ) + finally: + del llm + cleanup_dist_env_and_memory() + + def _assert_ar_runner(obj): from vllm.v1.worker.gpu.model_states.default import DefaultModelState from vllm.v1.worker.gpu.sample.sampler import Sampler diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index d062c282ae03..2a7c274e1139 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -555,6 +555,9 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: model_config.override_generation_config["max_new_tokens"] = None return + if vllm_config.parallel_config.pipeline_parallel_size != 1: + raise ValueError("Nemotron diffusion requires PP=1.") + attention_config = vllm_config.attention_config if attention_config.backend == AttentionBackendEnum.FLASHINFER: raise ValueError( diff --git a/vllm/model_executor/models/nemotron_dllm.py b/vllm/model_executor/models/nemotron_dllm.py index 8488df9a5250..0efecf0eb2ac 100644 --- a/vllm/model_executor/models/nemotron_dllm.py +++ b/vllm/model_executor/models/nemotron_dllm.py @@ -232,7 +232,7 @@ def _filter_sampling_logits( ) -> torch.Tensor: """Select sampling candidates after temperature scaling, excluding masks.""" logits = logits.float().clone() - logits[:, mask_token_id] = float("-inf") + logits[:, mask_token_id].fill_(-torch.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. @@ -300,6 +300,10 @@ def _compiled_masked_step( 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) + # Keep the random argmax separate from downstream gather/reductions. + # Inductor can otherwise emit an out-of-scope reduction temporary + # for the checkpoint vocabulary on SM100. + torch._dynamo.graph_break() else: x0 = logits_3d.argmax(dim=-1) # [num_decode, CL] chosen_logits = logits_3d.gather(-1, x0.unsqueeze(-1)).squeeze(-1) @@ -446,17 +450,20 @@ def __init__( 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 + if isinstance(slot_indices, int): + self.canvas[slot_indices : slot_indices + 1].fill_(self.mask_token_id) + else: + self.canvas.index_fill_(0, 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 + self.is_encoder_phase[slot_idx : slot_idx + 1].fill_(True) + self.init_canvas(slot_idx) + self.step[slot_idx : slot_idx + 1].zero_() + self.unmask_logprobs[slot_idx : slot_idx + 1].zero_() + self.unmask_ranks[slot_idx : slot_idx + 1].zero_() def remove_request(self, slot_idx: int) -> None: - self.is_encoder_phase[slot_idx] = False + self.is_encoder_phase[slot_idx : slot_idx + 1].fill_(False) class NemotronLabsDiffusionModelState(ModelState): @@ -592,7 +599,7 @@ def prepare_attn( slots ] if actual_num_reqs < num_reqs: - self._causal_buf[actual_num_reqs:num_reqs] = False + self._causal_buf[actual_num_reqs:num_reqs].fill_(False) causal: bool | torch.Tensor = self._causal_buf[:num_reqs] return build_attn_metadata( @@ -624,7 +631,7 @@ class MaskedDiffusionSampler: 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. + reveal time; top-k introspection uses that same distribution. """ def __init__( @@ -657,8 +664,7 @@ def __init__( 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. + # Per-slot top-k scores captured when each position is revealed. self._pending_logprobs: dict[int, LogprobsTensors] = {} def add_request(self, req_idx: int, sampling_params: Any) -> None: @@ -842,6 +848,9 @@ def __call__( # 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() + was_masked = (states.canvas[decode_slots] == self.mask_token_id) & ( + ~is_committing[:, None] + ) max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np) @@ -877,67 +886,70 @@ def __call__( 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. + # Capture each position's distribution at reveal, independently of + # neighbouring requests' logprob settings and convergence steps. 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], + revealed = was_masked & (states.canvas[decode_slots] != self.mask_token_id) + for li, slot in enumerate(decode_slots_np): + slot = int(slot) + requested = int(self.sampling_states.num_logprobs[slot]) + if requested <= 0: + continue + current = compute_topk_scores( + logits_3d[li], requested, states.canvas[slot] + ) + previous = self._pending_logprobs.get(slot) + if previous is None: + previous = LogprobsTensors( + logprob_token_ids=torch.zeros_like(current.logprob_token_ids), + logprobs=torch.full_like(current.logprobs, -torch.inf), + selected_token_ranks=torch.zeros_like( + current.selected_token_ranks + ), ) + write = revealed[li, :, None] + self._pending_logprobs[slot] = LogprobsTensors( + logprob_token_ids=torch.where( + write, current.logprob_token_ids, previous.logprob_token_ids + ), + logprobs=torch.where(write, current.logprobs, previous.logprobs), + selected_token_ranks=torch.where( + revealed[li], + current.selected_token_ranks, + previous.selected_token_ranks, + ), + ) - # 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. + # Fixed per-request offsets allow the output processor to select the + # emitted prefix using num_sampled without a host read of GPU phase flags. 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) - } + if max_num_logprobs >= 0 and num_decode > 0: 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] + for i, slot in enumerate(slots_np): + slot = int(slot) + lp = self._pending_logprobs.get(slot) + if self.sampling_states.num_logprobs[slot] > 0 and lp is not None: + token_ids, scores, ranks = ( + lp.logprob_token_ids, + lp.logprobs, + lp.selected_token_ranks, + ) 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, + token_ids = sampled[i].to(torch.int64).unsqueeze(-1) + scores = states.unmask_logprobs[slot].unsqueeze(-1) + ranks = states.unmask_ranks[slot] + padding = max_num_logprobs + 1 - scores.shape[1] + parts_ids.append(torch.nn.functional.pad(token_ids, (0, padding))) + parts_lp.append( + torch.nn.functional.pad(scores, (0, padding), value=-torch.inf) ) + parts_ranks.append(ranks) + 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=[i * CL for i in range(num_reqs)], + ) return self._build_output( input_batch,