@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / eb28452b10
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Model] Add DiffusionGemma Support (#45163)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Martin Kukla <martin.kukla@cantab.net> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Dipika Sikka <dsikka@redhat.com> Co-authored-by: NickLucche <nlucches@redhat.com> Co-authored-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: Alec Kohlhoff <134344302+aleckohlhoff@users.noreply.github.com> Co-authored-by: Porras Huang <20535584+porrashuang@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: scoootscooob <167050519+scoootscooob@users.noreply.github.com>

Lucas Wilkinson committed Jun 12, 2026 at 01:17 UTC eb28452b10a1376d143b2847a78b31726db346dd
52 files changed +2698 -235
benchmarks/kernels/benchmark_moe.py
+6
@@ -792,6 +792,12 @@ def get_model_params(config):
792 topk = text_config.num_experts_per_tok
793 intermediate_size = text_config.moe_intermediate_size
794 hidden_size = text_config.hidden_size
795 + elif architecture == "DiffusionGemmaForBlockDiffusion":
796 + text_config = config.get_text_config()
797 + E = text_config.num_experts
798 + topk = text_config.top_k_experts
799 + intermediate_size = text_config.moe_intermediate_size
800 + hidden_size = text_config.hidden_size
801 elif architecture == "HunYuanMoEV1ForCausalLM":
802 E = config.num_experts
803 topk = config.moe_topk[0]
cmake/external_projects/vllm_flash_attn.cmake
+1 -1
@@ -39,7 +39,7 @@ else()
39 FetchContent_Declare(
40 vllm-flash-attn
41 GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
42 - GIT_TAG dd62dac706b1cf7895bd99b18c6cb7e7e117ee25
42 + GIT_TAG 803020a8fa15407871341d41eba4919ade2ee1ee
43 GIT_PROGRESS TRUE
44 # Don't share the vllm-flash-attn build between build types
45 BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
docs/design/attention_backends.md
+1 -1
@@ -180,7 +180,7 @@ Priority is **1 = highest** (tried first).
180 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
181 | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
182 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
183 -| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
183 +| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |
184 | `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
185 | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
186
tests/kernels/attention/test_mixed_causal_attn.py new
+318
@@ -0,0 +1,318 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for per-request causal/non-causal attention (mixed batches).
4 +
5 +Validates that both triton and flash-attention backends correctly handle
6 +batches where some sequences use causal masking and others use non-causal
7 +(bidirectional) masking — needed by DiffusionGemma.
8 +"""
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm.platforms import current_platform
14 +from vllm.utils.torch_utils import set_random_seed
15 +
16 +# Mixed causal/non-causal attention is only validated on a subset of GPUs:
17 +# the Triton path on Hopper (SM90) and B200 (SM100); the FA4 path on Hopper
18 +# (SM90) only.
19 +_device_capability = current_platform.get_device_capability()
20 +_major = _device_capability.major if _device_capability is not None else None
21 +
22 +NUM_HEADS = [(4, 4), (8, 2)]
23 +HEAD_SIZES = [128]
24 +BLOCK_SIZES = [16]
25 +DTYPES = [torch.bfloat16]
26 +
27 +
28 +def ref_paged_attn(
29 + query: torch.Tensor,
30 + key_cache: torch.Tensor,
31 + value_cache: torch.Tensor,
32 + query_lens: list[int],
33 + kv_lens: list[int],
34 + block_tables: torch.Tensor,
35 + scale: float,
36 + per_seq_causal: list[bool],
37 + sliding_window: int | None = None,
38 +) -> torch.Tensor:
39 + num_seqs = len(query_lens)
40 + block_tables_np = block_tables.cpu().numpy()
41 + _, block_size, num_kv_heads, head_size = key_cache.shape
42 +
43 + outputs: list[torch.Tensor] = []
44 + start_idx = 0
45 + for i in range(num_seqs):
46 + query_len = query_lens[i]
47 + kv_len = kv_lens[i]
48 + q = query[start_idx : start_idx + query_len]
49 + q = q * scale
50 +
51 + num_kv_blocks = (kv_len + block_size - 1) // block_size
52 + block_indices = block_tables_np[i, :num_kv_blocks]
53 + k = key_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
54 + v = value_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
55 +
56 + if q.shape[1] != k.shape[1]:
57 + k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
58 + v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
59 +
60 + attn = torch.einsum("qhd,khd->hqk", q, k).float()
61 +
62 + if per_seq_causal[i]:
63 + mask = torch.triu(
64 + torch.ones(query_len, kv_len, device=attn.device),
65 + diagonal=kv_len - query_len + 1,
66 + ).bool()
67 + else:
68 + mask = torch.zeros(query_len, kv_len, device=attn.device).bool()
69 +
70 + if sliding_window is not None:
71 + sw_mask = (
72 + torch.triu(
73 + torch.ones(query_len, kv_len, device=attn.device),
74 + diagonal=kv_len - (query_len + sliding_window) + 1,
75 + )
76 + .bool()
77 + .logical_not()
78 + )
79 + mask |= sw_mask
80 +
81 + attn.masked_fill_(mask, float("-inf"))
82 + attn = torch.softmax(attn, dim=-1).to(v.dtype)
83 + out = torch.einsum("hqk,khd->qhd", attn, v)
84 + outputs.append(out)
85 + start_idx += query_len
86 +
87 + return torch.cat(outputs, dim=0)
88 +
89 +
90 +# ---- Triton backend test ----
91 +
92 +
93 +@pytest.mark.skipif(
94 + _major not in (9, 10),
95 + reason="Triton mixed causal attention requires Hopper (SM90) or B200 (SM100).",
96 +)
97 +@pytest.mark.parametrize(
98 + "seq_lens",
99 + [[(1, 128), (5, 64), (1, 256)]],
100 +)
101 +@pytest.mark.parametrize(
102 + "per_seq_causal",
103 + [[True, False, True], [False, True, False], [True, True, False]],
104 +)
105 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
106 +@pytest.mark.parametrize("head_size", HEAD_SIZES)
107 +@pytest.mark.parametrize("block_size", BLOCK_SIZES)
108 +@pytest.mark.parametrize("dtype", DTYPES)
109 +@torch.inference_mode()
110 +def test_triton_mixed_causal(
111 + seq_lens: list[tuple[int, int]],
112 + per_seq_causal: list[bool],
113 + num_heads: tuple[int, int],
114 + head_size: int,
115 + dtype: torch.dtype,
116 + block_size: int,
117 +):
118 + if not current_platform.is_cuda():
119 + pytest.skip("Triton attention requires CUDA")
120 +
121 + from vllm.v1.attention.ops.triton_unified_attention import unified_attention
122 +
123 + set_random_seed(42)
124 + device = "cuda"
125 +
126 + num_query_heads, num_kv_heads = num_heads
127 + assert len(seq_lens) == len(per_seq_causal)
128 +
129 + query_lens = [s[0] for s in seq_lens]
130 + kv_lens = [s[1] for s in seq_lens]
131 + num_seqs = len(seq_lens)
132 +
133 + num_query_tokens = sum(query_lens)
134 + max_kv_len = max(kv_lens)
135 + max_num_blocks = (max_kv_len + block_size - 1) // block_size
136 + num_blocks = max_num_blocks * num_seqs + 10
137 +
138 + scale = head_size**-0.5
139 + query = torch.randn(
140 + num_query_tokens, num_query_heads, head_size, dtype=dtype, device=device
141 + )
142 + key_cache = torch.randn(
143 + num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
144 + )
145 + value_cache = torch.randn(
146 + num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
147 + )
148 +
149 + block_tables_list = []
150 + for i in range(num_seqs):
151 + n_blocks = (kv_lens[i] + block_size - 1) // block_size
152 + blocks = list(range(i * max_num_blocks, i * max_num_blocks + n_blocks))
153 + blocks += [0] * (max_num_blocks - n_blocks)
154 + block_tables_list.append(blocks)
155 + block_tables = torch.tensor(block_tables_list, dtype=torch.int32, device=device)
156 +
157 + cu_seqlens_q = torch.zeros(num_seqs + 1, dtype=torch.int32, device=device)
158 + for i, ql in enumerate(query_lens):
159 + cu_seqlens_q[i + 1] = cu_seqlens_q[i] + ql
160 +
161 + seqused_k = torch.tensor(kv_lens, dtype=torch.int32, device=device)
162 + max_seqlen_q = max(query_lens)
163 + max_seqlen_k = max(kv_lens)
164 +
165 + causal_tensor = torch.tensor(per_seq_causal, dtype=torch.bool, device=device)
166 +
167 + output = torch.empty_like(query)
168 + unified_attention(
169 + q=query,
170 + k=key_cache,
171 + v=value_cache,
172 + out=output,
173 + cu_seqlens_q=cu_seqlens_q,
174 + max_seqlen_q=max_seqlen_q,
175 + seqused_k=seqused_k,
176 + max_seqlen_k=max_seqlen_k,
177 + softmax_scale=scale,
178 + causal=causal_tensor,
179 + window_size=(-1, -1),
180 + block_table=block_tables,
181 + softcap=0.0,
182 + q_descale=None,
183 + k_descale=1.0,
184 + v_descale=1.0,
185 + )
186 +
187 + ref_output = ref_paged_attn(
188 + query,
189 + key_cache,
190 + value_cache,
191 + query_lens,
192 + kv_lens,
193 + block_tables,
194 + scale,
195 + per_seq_causal,
196 + )
197 +
198 + torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
199 +
200 +
201 +# ---- Flash Attention 4 backend test (native per_seq_causal) ----
202 +
203 +
204 +@pytest.mark.skipif(
205 + _major != 9,
206 + reason="FA4 mixed causal attention requires Hopper (SM90).",
207 +)
208 +@pytest.mark.parametrize(
209 + "seq_lens",
210 + [[(1, 128), (5, 64), (1, 256)]],
211 +)
212 +@pytest.mark.parametrize(
213 + "per_seq_causal",
214 + [[True, False, True], [False, True, False]],
215 +)
216 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
217 +@pytest.mark.parametrize("head_size", HEAD_SIZES)
218 +@pytest.mark.parametrize("block_size", BLOCK_SIZES)
219 +@pytest.mark.parametrize("dtype", DTYPES)
220 +@torch.inference_mode()
221 +def test_flash_attn4_mixed_causal(
222 + seq_lens: list[tuple[int, int]],
223 + per_seq_causal: list[bool],
224 + num_heads: tuple[int, int],
225 + head_size: int,
226 + dtype: torch.dtype,
227 + block_size: int,
228 +):
229 + if not current_platform.is_cuda():
230 + pytest.skip("Flash attention requires CUDA")
231 +
232 + try:
233 + from vllm.vllm_flash_attn import (
234 + fa_version_unsupported_reason,
235 + flash_attn_varlen_func,
236 + is_fa_version_supported,
237 + )
238 + except ImportError:
239 + pytest.skip("vllm_flash_attn not available")
240 +
241 + if not is_fa_version_supported(4):
242 + reason = fa_version_unsupported_reason(4)
243 + pytest.skip(f"FA4 not supported: {reason}")
244 +
245 + set_random_seed(42)
246 + device = "cuda"
247 +
248 + num_query_heads, num_kv_heads = num_heads
249 + assert len(seq_lens) == len(per_seq_causal)
250 +
251 + query_lens = [s[0] for s in seq_lens]
252 + kv_lens = [s[1] for s in seq_lens]
253 + num_seqs = len(seq_lens)
254 +
255 + num_query_tokens = sum(query_lens)
256 + max_kv_len = max(kv_lens)
257 + max_num_blocks = (max_kv_len + block_size - 1) // block_size
258 + num_blocks = max_num_blocks * num_seqs + 10
259 +
260 + scale = head_size**-0.5
261 + query = torch.randn(
262 + num_query_tokens, num_query_heads, head_size, dtype=dtype, device=device
263 + )
264 + key_cache = torch.randn(
265 + num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
266 + )
267 + value_cache = torch.randn(
268 + num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
269 + )
270 +
271 + block_tables_list = []
272 + for i in range(num_seqs):
273 + n_blocks = (kv_lens[i] + block_size - 1) // block_size
274 + blocks = list(range(i * max_num_blocks, i * max_num_blocks + n_blocks))
275 + blocks += [0] * (max_num_blocks - n_blocks)
276 + block_tables_list.append(blocks)
277 + block_tables = torch.tensor(block_tables_list, dtype=torch.int32, device=device)
278 +
279 + cu_seqlens_q = torch.zeros(num_seqs + 1, dtype=torch.int32, device=device)
280 + for i, ql in enumerate(query_lens):
281 + cu_seqlens_q[i + 1] = cu_seqlens_q[i] + ql
282 +
283 + seqused_k = torch.tensor(kv_lens, dtype=torch.int32, device=device)
284 + per_seq_causal_tensor = torch.tensor(
285 + per_seq_causal, dtype=torch.int32, device=device
286 + )
287 +
288 + ref_output = ref_paged_attn(
289 + query,
290 + key_cache,
291 + value_cache,
292 + query_lens,
293 + kv_lens,
294 + block_tables,
295 + scale,
296 + per_seq_causal,
297 + )
298 +
299 + output = torch.empty_like(query)
300 + flash_attn_varlen_func(
301 + q=query,
302 + k=key_cache,
303 + v=value_cache,
304 + out=output,
305 + cu_seqlens_q=cu_seqlens_q,
306 + max_seqlen_q=max(query_lens),
307 + seqused_k=seqused_k,
308 + max_seqlen_k=max(kv_lens),
309 + softmax_scale=scale,
310 + # The kernel must be compiled causal for `dynamic_causal` to take effect.
311 + causal=True,
312 + block_table=block_tables,
313 + softcap=0.0,
314 + dynamic_causal=per_seq_causal_tensor,
315 + fa_version=4,
316 + )
317 +
318 + torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
tests/models/registry.py
+4
@@ -898,6 +898,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
898 ),
899 "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"),
900 "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"),
901 + "DiffusionGemmaForBlockDiffusion": _HfExamplesInfo(
902 + "google/diffusiongemma-26B-A4B-it",
903 + trust_remote_code=True,
904 + ),
905 "Gemma4ForConditionalGeneration": _HfExamplesInfo(
906 "google/gemma-4-E2B-it",
907 min_transformers_version="5.5.0",
tests/models/utils.py
+3 -1
@@ -486,6 +486,7 @@ def dummy_hf_overrides(
486 "Gemma3nForConditionalGeneration",
487 "Gemma4ForCausalLM",
488 "Gemma4ForConditionalGeneration",
489 + "DiffusionGemmaForBlockDiffusion",
490 )
491 else 1
492 )
@@ -558,7 +559,8 @@ def dummy_hf_overrides(
559 )
560
561 # e.g.: Qwen/Qwen2-Audio-7B-Instruct
561 - if hasattr(hf_config, "audio_config"):
562 + # audio_config may exist but be None (e.g. audio-less Gemma4 variants).
563 + if getattr(hf_config, "audio_config", None) is not None:
564 hf_config.audio_config.update(
565 {
566 "num_layers": 1,
tests/tool_parsers/test_gemma4_tool_parser.py
+82
@@ -702,6 +702,88 @@ class TestStreamingExtraction:
702 ' <meta name="viewport" content="width=device-width">\n'
703 )
704
705 + def _collect_tool_calls_by_index(self, results):
706 + """Group streamed tool-call fragments by their ``index``.
707 +
708 + Returns ``{index: {"name": str | None, "arguments": str}}`` where
709 + ``arguments`` is the concatenation of every streamed argument
710 + fragment for that index (which should form valid JSON once complete).
711 + """
712 + by_index: dict[int, dict[str, Any]] = {}
713 + for delta, _ in results:
714 + if not (delta and delta.tool_calls):
715 + continue
716 + for tc in delta.tool_calls:
717 + entry = by_index.setdefault(tc.index, {"name": None, "arguments": ""})
718 + func = tc.function
719 + if isinstance(func, dict):
720 + name = func.get("name")
721 + arg = func.get("arguments", "")
722 + else:
723 + name = getattr(func, "name", None)
724 + arg = getattr(func, "arguments", "") or ""
725 + if name:
726 + entry["name"] = name
727 + if arg:
728 + entry["arguments"] += arg
729 + return by_index
730 +
731 + def test_streaming_single_chunk_complete_tool_call(self, parser, mock_request):
732 + """A backend may deliver a whole tool call in one streaming delta.
733 +
734 + The start token, ``call:name{...}`` payload and the end token all
735 + arrive in a single chunk. The parser must still emit one
736 + ``DeltaToolCall`` with the correct name + complete arguments JSON
737 + (rather than swallowing it and finishing with finish_reason="stop").
738 + """
739 + chunks = [
740 + '<|tool_call>call:name_a_color{color_hex:<|"|>00ff11<|"|>}<tool_call|>',
741 + ]
742 +
743 + results = self._simulate_streaming(parser, mock_request, chunks)
744 +
745 + # Exactly one delta should carry tool_calls, and it must not be
746 + # emitted as plain content (which would yield finish_reason="stop").
747 + tool_call_deltas = [
748 + delta for delta, _ in results if delta is not None and delta.tool_calls
749 + ]
750 + assert len(tool_call_deltas) == 1, (
751 + "Expected exactly one delta carrying the batched tool call"
752 + )
753 + assert all(
754 + delta.content is None for delta, _ in results if delta is not None
755 + ), "Complete tool call must not leak as content"
756 +
757 + by_index = self._collect_tool_calls_by_index(results)
758 + assert set(by_index) == {0}
759 + assert by_index[0]["name"] == "name_a_color"
760 + assert json.loads(by_index[0]["arguments"]) == {"color_hex": "00ff11"}
761 +
762 + def test_streaming_multi_chunk_batched_tool_calls(self, parser, mock_request):
763 + """A single delta may batch MULTIPLE complete tool calls.
764 +
765 + ``<|tool_call>...<tool_call|><|tool_call>...<tool_call|>`` arriving in
766 + one chunk must emit BOTH calls (one DeltaToolCall each, with distinct
767 + indices), not just the first.
768 + """
769 + chunks = [
770 + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}<tool_call|>'
771 + '<|tool_call>call:get_time{timezone:<|"|>GMT<|"|>}<tool_call|>',
772 + ]
773 +
774 + results = self._simulate_streaming(parser, mock_request, chunks)
775 +
776 + by_index = self._collect_tool_calls_by_index(results)
777 + assert set(by_index) == {0, 1}, (
778 + f"Expected two tool calls (indices 0 and 1), got {sorted(by_index)}"
779 + )
780 +
781 + assert by_index[0]["name"] == "get_weather"
782 + assert json.loads(by_index[0]["arguments"]) == {"location": "London"}
783 +
784 + assert by_index[1]["name"] == "get_time"
785 + assert json.loads(by_index[1]["arguments"]) == {"timezone": "GMT"}
786 +
787 def test_streaming_trailing_bare_bool_not_duplicated(self, parser, mock_request):
788 """Trailing bare boolean must not be streamed twice."""
789 chunks = [
tests/v1/cudagraph/test_cudagraph_dispatch.py
+1
@@ -49,6 +49,7 @@ def _create_vllm_config(
49 )
50 mock_config.parallel_config = ParallelConfig()
51 mock_config.speculative_config = None # No speculative decoding
52 + mock_config.num_speculative_tokens = 0
53 if not lora_config:
54 mock_config.lora_config = None
55 else:
tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py
+1 -1
@@ -104,7 +104,7 @@ def _run_engine_core_handshake(
104 speculative_config=None,
105 ec_transfer_config=None,
106 max_concurrent_batches=1,
107 - model_config=SimpleNamespace(runner_type="generate"),
107 + model_config=SimpleNamespace(runner_type="generate", is_diffusion=False),
108 cache_config=SimpleNamespace(
109 enable_prefix_caching=False,
110 prefix_caching_hash_algo="builtin",
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+14 -5
@@ -70,6 +70,7 @@ def _make_runner(**overrides: Any) -> Any:
70 runner.use_aux_hidden_state_outputs = False
71 runner.speculative_config = None
72 runner.speculator = None
73 + runner.num_speculative_steps = 0
74 runner.encoder_cache = None
75 runner.is_pooling_model = False
76 runner.is_last_pp_rank = True
@@ -102,18 +103,22 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch):
103 lambda load_config: SimpleNamespace(load_model=lambda **_: model),
104 )
105 monkeypatch.setattr(mrv2, "prepare_communication_buffer_for_model", prepared.append)
105 - monkeypatch.setattr(mrv2, "init_model_state", lambda *args: "model-state")
106 + monkeypatch.setattr(
107 + mrv2,
108 + "init_model_state",
109 + lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1),
110 + )
111 monkeypatch.setattr(
112 eplb,
113 "is_mixture_of_experts",
114 lambda loaded_model: getattr(loaded_model, "is_moe", False),
115 )
116
112 - runner = _make_runner()
117 + runner = _make_runner(is_last_pp_rank=False)
118 mrv2.GPUModelRunner.load_model(runner)
119
120 assert runner.model is model
116 - assert runner.model_state == "model-state"
121 + assert runner.model_state is not None
122 assert prepared == [model]
123 assert runner.eplb_state is not None
124 assert runner.eplb_state.add_model_calls == [(model, runner.model_config)]
@@ -133,10 +138,14 @@ def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch):
138 lambda load_config: SimpleNamespace(load_model=lambda **_: model),
139 )
140 monkeypatch.setattr(mrv2, "prepare_communication_buffer_for_model", prepared.append)
136 - monkeypatch.setattr(mrv2, "init_model_state", lambda *args: "model-state")
141 + monkeypatch.setattr(
142 + mrv2,
143 + "init_model_state",
144 + lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1),
145 + )
146 monkeypatch.setattr(eplb, "is_mixture_of_experts", lambda *_: True)
147
139 - runner = _make_runner()
148 + runner = _make_runner(is_last_pp_rank=False)
149 mrv2.GPUModelRunner.load_model(runner, load_dummy_weights=True)
150
151 assert runner.load_config.load_format == "dummy"
vllm/benchmarks/serve.py
+117 -1
@@ -248,6 +248,68 @@ async def fetch_spec_decode_metrics(
248 return None
249
250
251 +@dataclass
252 +class DiffusionMetrics:
253 + """Diffusion (dLLM) decoding metrics from the server's Prometheus endpoint."""
254 +
255 + num_denoising_steps: int
256 + num_canvas_positions: int
257 + num_committed_tokens: int
258 +
259 +
260 +async def fetch_diffusion_metrics(
261 + base_url: str, session: aiohttp.ClientSession
262 +) -> DiffusionMetrics | None:
263 + """Fetch diffusion decoding metrics from the server's Prometheus endpoint.
264 +
265 + Returns None if the model is not a diffusion model or metrics are not
266 + available.
267 + """
268 + metrics_url = f"{base_url}/metrics"
269 + try:
270 + async with session.get(metrics_url) as response:
271 + if response.status != 200:
272 + return None
273 + text = await response.text()
274 +
275 + num_denoising_steps = 0
276 + num_canvas_positions = 0
277 + num_committed_tokens = 0
278 + found_diffusion = False
279 +
280 + for line in text.split("\n"):
281 + line = line.strip()
282 + if not line or line.startswith("#"):
283 + continue
284 +
285 + if line.startswith("vllm:diffusion"):
286 + # Extract metric name (before labels) to avoid matching
287 + # substrings inside label values.
288 + parts = line.split(None, 1)
289 + metric_name = parts[0].split("{")[0]
290 + if not metric_name.endswith("_total"):
291 + continue
292 + found_diffusion = True
293 + with contextlib.suppress(ValueError):
294 + if "num_denoising_steps" in metric_name:
295 + num_denoising_steps += int(float(parts[-1]))
296 + elif "num_canvas_positions" in metric_name:
297 + num_canvas_positions += int(float(parts[-1]))
298 + elif "num_committed_tokens" in metric_name:
299 + num_committed_tokens += int(float(parts[-1]))
300 +
301 + if not found_diffusion:
302 + return None
303 +
304 + return DiffusionMetrics(
305 + num_denoising_steps=num_denoising_steps,
306 + num_canvas_positions=num_canvas_positions,
307 + num_committed_tokens=num_committed_tokens,
308 + )
309 + except (aiohttp.ClientError, asyncio.TimeoutError):
310 + return None
311 +
312 +
313 class TaskType(Enum):
314 GENERATION = "generation"
315 POOLING = "pooling"
@@ -887,6 +949,7 @@ async def benchmark(
949 print("Self timing is set, using the timestamps from the trace file.")
950
951 spec_decode_metrics_before = await fetch_spec_decode_metrics(base_url, session)
952 + diffusion_metrics_before = await fetch_diffusion_metrics(base_url, session)
953
954 pbar = None if disable_tqdm else tqdm(total=len(input_requests))
955
@@ -1016,6 +1079,34 @@ async def benchmark(
1079 "per_position_acceptance_rates": per_pos_rates,
1080 }
1081
1082 + diffusion_metrics_after = await fetch_diffusion_metrics(base_url, session)
1083 + diffusion_stats: dict[str, Any] | None = None
1084 + if diffusion_metrics_before is not None and diffusion_metrics_after is not None:
1085 + delta_steps = (
1086 + diffusion_metrics_after.num_denoising_steps
1087 + - diffusion_metrics_before.num_denoising_steps
1088 + )
1089 + delta_positions = (
1090 + diffusion_metrics_after.num_canvas_positions
1091 + - diffusion_metrics_before.num_canvas_positions
1092 + )
1093 + delta_committed = (
1094 + diffusion_metrics_after.num_committed_tokens
1095 + - diffusion_metrics_before.num_committed_tokens
1096 + )
1097 + if delta_steps > 0 and delta_committed > 0:
1098 + block_size = delta_positions / delta_steps # canvas length (CL)
1099 + num_canvases = delta_committed / block_size # = number of commit steps
1100 + denoising_steps = delta_steps - num_canvases # exclude commit steps
1101 + diffusion_stats = {
1102 + "denoising_steps": denoising_steps,
1103 + "canvas_positions": delta_positions,
1104 + "committed_tokens": delta_committed,
1105 + "committed_throughput": delta_committed / benchmark_duration,
1106 + "steps_per_canvas": denoising_steps / num_canvases,
1107 + "committed_per_step": delta_committed / denoising_steps,
1108 + }
1109 +
1110 if task_type == TaskType.GENERATION:
1111 metrics, actual_output_lens = calculate_metrics(
1112 input_requests=input_requests,
@@ -1134,6 +1225,16 @@ async def benchmark(
1225 "per_position_acceptance_rates", []
1226 )
1227
1228 + if diffusion_stats is not None:
1229 + result["diffusion_committed_throughput"] = diffusion_stats[
1230 + "committed_throughput"
1231 + ]
1232 + result["diffusion_steps_per_canvas"] = diffusion_stats["steps_per_canvas"]
1233 + result["diffusion_committed_per_step"] = diffusion_stats["committed_per_step"]
1234 + result["diffusion_committed_tokens"] = int(diffusion_stats["committed_tokens"])
1235 + result["diffusion_denoising_steps"] = int(diffusion_stats["denoising_steps"])
1236 + result["diffusion_canvas_positions"] = int(diffusion_stats["canvas_positions"])
1237 +
1238 def process_one_metric(
1239 # E.g., "ttft"
1240 metric_attribute_name: str,
@@ -1179,7 +1280,22 @@ async def benchmark(
1280 process_one_metric("itl", "ITL", "Inter-token Latency")
1281 process_one_metric("e2el", "E2EL", "End-to-end Latency")
1282
1182 - if spec_decode_stats is not None:
1283 + if diffusion_stats is not None:
1284 + print("{s:{c}^{n}}".format(s="Diffusion Decoding", n=50, c="-"))
1285 + for label, key, value_fmt in (
1286 + ("Committed throughput (tok/s):", "committed_throughput", "{:<10.2f}"),
1287 + ("Denoising steps per canvas:", "steps_per_canvas", "{:<10.2f}"),
1288 + ("Committed per denoising step:", "committed_per_step", "{:<10.2f}"),
1289 + ("Committed tokens:", "committed_tokens", "{:<10d}"),
1290 + ("Denoising steps:", "denoising_steps", "{:<10d}"),
1291 + ("Canvas positions evaluated:", "canvas_positions", "{:<10d}"),
1292 + ):
1293 + value = diffusion_stats[key]
1294 + if value_fmt.endswith("d}"):
1295 + value = int(value)
1296 + print("{:<40} ".format(label) + value_fmt.format(value))
1297 +
1298 + if spec_decode_stats is not None and diffusion_stats is None:
1299 print("{s:{c}^{n}}".format(s="Speculative Decoding", n=50, c="-"))
1300 print(
1301 "{:<40} {:<10.2f}".format(
vllm/config/__init__.py
+3
@@ -10,6 +10,7 @@ from vllm.config.compilation import (
10 PassConfig,
11 )
12 from vllm.config.device import DeviceConfig
13 +from vllm.config.diffusion import DiffusionConfig
14 from vllm.config.ec_transfer import ECTransferConfig
15 from vllm.config.kernel import KernelConfig
16 from vllm.config.kv_events import KVEventsConfig
@@ -72,6 +73,8 @@ __all__ = [
73 "PassConfig",
74 # From vllm.config.device
75 "DeviceConfig",
76 + # From vllm.config.diffusion
77 + "DiffusionConfig",
78 # From vllm.config.ec_transfer
79 "ECTransferConfig",
80 # From vllm.config.kernel
vllm/config/diffusion.py new
+26
@@ -0,0 +1,26 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Configuration for discrete diffusion (dLLM) models."""
4 +
5 +from pydantic import Field
6 +
7 +from vllm.config.utils import config
8 +
9 +
10 +@config
11 +class DiffusionConfig:
12 + """Configuration for discrete diffusion language models (dLLMs).
13 +
14 + dLLMs generate tokens via iterative denoising over a fixed-length canvas
15 + rather than left-to-right autoregressive decoding. They reuse the
16 + speculative-decoding data path (draft token ids, scheduled spec decode
17 + tokens) with overloaded semantics for block-based generation.
18 + """
19 +
20 + canvas_length: int = Field(default=None, gt=0) # type: ignore[assignment]
21 + """Length of the denoising canvas (block). Also determines the number of
22 + speculative tokens scheduled per step."""
23 +
24 + max_denoising_steps: int | None = None
25 + """Maximum number of denoising iterations per canvas block.
26 + If not set, read from the model's generation_config.json."""
vllm/config/model.py
+5
@@ -1546,6 +1546,11 @@ class ModelConfig:
1546 """Extract the HF encoder/decoder model flag."""
1547 return is_encoder_decoder(self.hf_config)
1548
1549 + @cached_property
1550 + def is_diffusion(self) -> bool:
1551 + """Detect discrete diffusion (dLLM) models from HF config."""
1552 + return getattr(self.hf_config, "canvas_length", None) is not None
1553 +
1554 @property
1555 def uses_alibi(self) -> bool:
1556 cfg = self.hf_text_config
vllm/config/vllm.py
+13 -6
@@ -31,6 +31,7 @@ from .attention import AttentionConfig
31 from .cache import CacheConfig
32 from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode
33 from .device import DeviceConfig
34 +from .diffusion import DiffusionConfig
35 from .ec_transfer import ECTransferConfig
36 from .kernel import KernelConfig
37 from .kv_events import KVEventsConfig
@@ -323,6 +324,9 @@ class VllmConfig:
324 """LoRA configuration."""
325 speculative_config: SpeculativeConfig | None = None
326 """Speculative decoding configuration."""
327 + diffusion_config: DiffusionConfig | None = None
328 + """Diffusion LLM (dLLM) configuration."""
329 +
330 structured_outputs_config: StructuredOutputsConfig = Field(
331 default_factory=StructuredOutputsConfig
332 )
@@ -511,6 +515,11 @@ class VllmConfig:
515 and self.speculative_config.num_speculative_tokens is not None
516 ):
517 return self.speculative_config.num_speculative_tokens
518 + if (
519 + self.diffusion_config is not None
520 + and self.diffusion_config.canvas_length is not None
521 + ):
522 + return self.diffusion_config.canvas_length
523 return 0
524
525 @property
@@ -519,6 +528,9 @@ class VllmConfig:
528 if use_v2_model_runner is not None:
529 return use_v2_model_runner
530
531 + if self.model_config is not None and self.model_config.is_diffusion:
532 + return True
533 +
534 if not self._is_default_v2_model_runner_model():
535 return False
536
@@ -1654,12 +1666,7 @@ class VllmConfig:
1666 self.compilation_config.max_cudagraph_capture_size
1667 )
1668 if max_cudagraph_capture_size is None:
1657 - decode_query_len = 1
1658 - if (
1659 - self.speculative_config
1660 - and self.speculative_config.num_speculative_tokens
1661 - ):
1662 - decode_query_len += self.speculative_config.num_speculative_tokens
1669 + decode_query_len = 1 + self.num_speculative_tokens
1670 max_cudagraph_capture_size = min(
1671 self.scheduler_config.max_num_seqs * decode_query_len * 2, 512
1672 )
vllm/engine/arg_utils.py
+16
@@ -38,6 +38,7 @@ from vllm.config import (
38 CompilationConfig,
39 ConfigType,
40 DeviceConfig,
41 + DiffusionConfig,
42 ECTransferConfig,
43 EPLBConfig,
44 KernelConfig,
@@ -616,6 +617,7 @@ class EngineArgs:
617 spec_method: str | None = None
618 spec_model: str | None = None
619 spec_tokens: int | None = None
620 + diffusion_config: dict[str, Any] | None = None
621
622 show_hidden_metrics_for_version: str | None = (
623 ObservabilityConfig.show_hidden_metrics_for_version
@@ -1473,6 +1475,10 @@ class EngineArgs:
1475 vllm_group.add_argument(
1476 "--spec-tokens", **speculative_kwargs["num_speculative_tokens"]
1477 )
1478 + vllm_kwargs["diffusion_config"]["type"] = optional_type(json.loads)
1479 + vllm_group.add_argument(
1480 + "--diffusion-config", "-dc", **vllm_kwargs["diffusion_config"]
1481 + )
1482 vllm_group.add_argument(
1483 "--kv-transfer-config", **vllm_kwargs["kv_transfer_config"]
1484 )
@@ -1702,6 +1708,14 @@ class EngineArgs:
1708 )
1709 return SpeculativeConfig(**self.speculative_config)
1710
1711 + def create_diffusion_config(self) -> DiffusionConfig | None:
1712 + if self.diffusion_config is None:
1713 + return None
1714 + cfg = self.diffusion_config
1715 + if isinstance(cfg, str):
1716 + cfg = json.loads(cfg)
1717 + return DiffusionConfig(**cfg)
1718 +
1719 def create_engine_config(
1720 self,
1721 usage_context: UsageContext | None = None,
@@ -2016,6 +2030,7 @@ class EngineArgs:
2030 target_model_config=model_config,
2031 target_parallel_config=parallel_config,
2032 )
2033 + diffusion_config = self.create_diffusion_config()
2034
2035 self._set_default_max_num_seqs_and_batched_tokens_args(
2036 usage_context,
@@ -2243,6 +2258,7 @@ class EngineArgs:
2258 kernel_config=kernel_config,
2259 lora_config=lora_config,
2260 speculative_config=speculative_config,
2261 + diffusion_config=diffusion_config,
2262 structured_outputs_config=self.structured_outputs_config,
2263 observability_config=observability_config,
2264 compilation_config=compilation_config,
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
+2
@@ -188,6 +188,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
188 def _supports_activation(activation: MoEActivation) -> bool:
189 return activation in [
190 MoEActivation.SILU,
191 + MoEActivation.GELU_TANH,
192 MoEActivation.RELU2_NO_MUL,
193 MoEActivation.SWIGLUOAI,
194 ]
@@ -267,6 +268,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
268
269 activation_str_to_value_map = {
270 MoEActivation.SILU: ActivationType.Swiglu, # This is the default
271 + MoEActivation.GELU_TANH: ActivationType.Geglu,
272 MoEActivation.SWIGLUOAI: ActivationType.Swiglu, # gpt-oss alias
273 MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
274 }
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+1
@@ -142,6 +142,7 @@ class TrtLlmNvFp4ExpertsBase:
142 MoEActivation.SILU,
143 MoEActivation.RELU2_NO_MUL,
144 MoEActivation.GELU,
145 + MoEActivation.GELU_TANH,
146 ]
147
148 @staticmethod
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+1
@@ -34,6 +34,7 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType"
34 MoEActivation.GELU_NO_MUL: ActivationType.Gelu,
35 MoEActivation.SILU: ActivationType.Swiglu,
36 MoEActivation.GELU: ActivationType.Geglu,
37 + MoEActivation.GELU_TANH: ActivationType.Geglu,
38 MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
39 }
40 return ACTIVATION_TO_FI_ACTIVATION[activation]
vllm/model_executor/models/config.py
+55
@@ -105,6 +105,60 @@ class Gemma4Config(VerifyAndUpdateConfig):
105 )
106
107
108 +class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig):
109 + @classmethod
110 + def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
111 + """Set up the diffusion config and defaults for DiffusionGemma.
112 +
113 + Auto-creates DiffusionConfig from the HF config when the user
114 + didn't pass ``--diffusion-config``. Diffusion sampling params are
115 + read straight from generation_config.json at sampler-build time
116 + (see DiffusionGemma's custom_sampler), not injected here.
117 + """
118 + # Inherit Gemma4's attention backend selection (FA4 on Hopper,
119 + # TRITON_ATTN fallback for heterogeneous head dims).
120 + Gemma4Config.verify_and_update_config(vllm_config)
121 +
122 + from vllm.v1.attention.backends.registry import AttentionBackendEnum
123 +
124 + attention_config = vllm_config.attention_config
125 + if attention_config.backend == AttentionBackendEnum.FLASHINFER:
126 + raise ValueError(
127 + "FlashInfer does not support DiffusionGemma's mixed "
128 + "causal/bidirectional attention. Use --attention-backend "
129 + "FLASH_ATTN or TRITON_ATTN instead."
130 + )
131 + if attention_config.backend is None and not attention_config.use_non_causal:
132 + attention_config.use_non_causal = True
133 + logger.info(
134 + "DiffusionGemma uses mixed causal/bidirectional attention "
135 + "within a batch; setting use_non_causal=True to exclude "
136 + "FlashInfer from auto-selection."
137 + )
138 +
139 + # Auto-create DiffusionConfig from HF config if not provided.
140 + if vllm_config.diffusion_config is None:
141 + from vllm.config.diffusion import DiffusionConfig
142 +
143 + hf_config = vllm_config.model_config.hf_config
144 + canvas_length = getattr(hf_config, "canvas_length", 256)
145 + vllm_config.diffusion_config = DiffusionConfig(
146 + canvas_length=canvas_length,
147 + )
148 +
149 + # The diffusion sampler materializes [num_seqs, canvas_length, vocab]
150 + # fp32 transients, so concurrency is memory-bound (>8 OOMs a single H200).
151 + # Default to 8 when the user didn't pass --max-num-seqs.
152 + # We can't see the original None here (the engine already filled a generic
153 + # default), so use >= DEFAULT_MAX_NUM_SEQS as a proxy, (the default is much
154 + # larger than any deliberate value for this model)
155 + from vllm.config.scheduler import SchedulerConfig
156 +
157 + sc = vllm_config.scheduler_config
158 + if sc is not None and sc.max_num_seqs >= SchedulerConfig.DEFAULT_MAX_NUM_SEQS:
159 + sc.max_num_seqs = 8
160 +
161 +
162 class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig):
163 @staticmethod
164 def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -591,6 +645,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
645 "ColQwen3_5": Qwen3_5ForConditionalGenerationConfig,
646 "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig,
647 "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
648 + "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501
649 "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501
650 "FalconMambaForCausalLM": MambaModelConfig,
651 "Gemma3TextModel": Gemma3TextModelConfig,
vllm/model_executor/models/diffusion_gemma.py new
+1363
@@ -0,0 +1,1363 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DiffusionGemma model, ModelState, and Sampler for vLLM.
4 +
5 +Single Gemma4 backbone run in two modes (like YOCO):
6 +- encoder mode: causal attention, writes KV cache
7 +- decoder mode: bidirectional attention, reads encoder KV, doesn't write
8 +
9 +Same weights, same layers. The only decoder-unique component is a
10 +self-conditioning MLP.
11 +
12 +Multimodal support: the model always includes a vision tower (shared with Gemma4).
13 +Images are encoded through the vision tower and projected into the LM embedding space
14 +via Gemma4MultimodalEmbedder.
15 +"""
16 +
17 +from __future__ import annotations
18 +
19 +from collections.abc import Iterable, Mapping
20 +from types import SimpleNamespace
21 +from typing import Any
22 +
23 +import numpy as np
24 +import torch
25 +from torch import nn
26 +from torch.nn import functional as F
27 +from transformers import AutoModel
28 +
29 +from vllm.config import VllmConfig
30 +from vllm.config.compilation import CUDAGraphMode
31 +from vllm.logger import init_logger
32 +from vllm.model_executor.layers.layernorm import RMSNorm
33 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
34 +from vllm.model_executor.layers.vocab_parallel_embedding import (
35 + ParallelLMHead,
36 +)
37 +from vllm.model_executor.models.gemma4 import Gemma4Model
38 +from vllm.model_executor.models.gemma4_mm import (
39 + Gemma4DummyInputsBuilder,
40 + Gemma4ForConditionalGeneration,
41 + Gemma4MultimodalEmbedder,
42 + Gemma4MultiModalProcessor,
43 + Gemma4ProcessingInfo,
44 +)
45 +from vllm.model_executor.models.module_mapping import MultiModelKeys
46 +from vllm.model_executor.models.transformers.utils import recursive_replace_linear
47 +from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
48 +from vllm.multimodal import MULTIMODAL_REGISTRY
49 +from vllm.v1.outputs import LogprobsTensors
50 +from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
51 +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu
52 +from vllm.v1.worker.gpu.model_states.interface import ModelState
53 +from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
54 +from vllm.v1.worker.gpu.sample.output import SamplerOutput
55 +from vllm.v1.worker.gpu.sample.penalties import use_penalty
56 +
57 +from .interfaces import (
58 + SupportsMultiModal,
59 + SupportsPP,
60 + SupportsQuant,
61 +)
62 +
63 +logger = init_logger(__name__)
64 +
65 +
66 +class DiffusionGemmaSelfConditioning(nn.Module):
67 + """Gated MLP that processes soft embeddings from the previous denoising step.
68 +
69 + Structurally identical to Gemma4MLP but with self_conditioning_size
70 + and post_norm without learned scale.
71 + """
72 +
73 + def __init__(
74 + self, hidden_size: int, self_conditioning_size: int, eps: float = 1e-6
75 + ):
76 + super().__init__()
77 + self.pre_norm = RMSNorm(hidden_size, eps=eps)
78 + self.post_norm = RMSNorm(hidden_size, eps=eps, has_weight=False)
79 + self.gate_proj = nn.Linear(hidden_size, self_conditioning_size, bias=False)
80 + self.up_proj = nn.Linear(hidden_size, self_conditioning_size, bias=False)
81 + self.down_proj = nn.Linear(self_conditioning_size, hidden_size, bias=False)
82 +
83 + def forward(
84 + self,
85 + inputs_embeds: torch.Tensor,
86 + soft_embeds: torch.Tensor,
87 + ) -> torch.Tensor:
88 + x = self.pre_norm(soft_embeds)
89 + sc_signal = self.down_proj(
90 + F.gelu(self.gate_proj(x), approximate="tanh") * self.up_proj(x)
91 + )
92 + return self.post_norm(inputs_embeds + sc_signal)
93 +
94 +
95 +# ---------------------------------------------------------------------------
96 +# Multimodal processing info (overrides Gemma4 config type check)
97 +# ---------------------------------------------------------------------------
98 +
99 +
100 +class DiffusionGemmaProcessingInfo(Gemma4ProcessingInfo):
101 + """Processing info for DiffusionGemma.
102 +
103 + Overrides ``get_hf_config`` to accept ``DiffusionGemmaConfig``
104 + (which inherits from ``PretrainedConfig``, not ``Gemma4Config``).
105 + Supports image and video modalities.
106 + """
107 +
108 + def get_hf_config(self):
109 + # DiffusionGemmaConfig doesn't inherit from Gemma4Config, so we
110 + # accept any PretrainedConfig here.
111 + return self.ctx.get_hf_config()
112 +
113 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
114 + # DiffusionGemma supports image and video inputs.
115 + return {"image": None, "video": None}
116 +
117 + def get_mm_max_tokens_per_item(
118 + self, seq_len: int, mm_counts: Mapping[str, int]
119 + ) -> Mapping[str, int] | None:
120 + return super().get_mm_max_tokens_per_item(seq_len, mm_counts)
121 +
122 +
123 +@torch.compile(dynamic=True)
124 +def _softcap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor:
125 + # fp32 before tanh for numerical stability (matches HF DiffusionGemma).
126 + # Compiling fuses the cast/div/tanh/mul into one elementwise kernel over
127 + # the [num_tokens, vocab] logits instead of four separate passes.
128 + logits = logits.float()
129 + return torch.tanh(logits / cap) * cap
130 +
131 +
132 +@MULTIMODAL_REGISTRY.register_processor(
133 + Gemma4MultiModalProcessor,
134 + info=DiffusionGemmaProcessingInfo,
135 + dummy_inputs=Gemma4DummyInputsBuilder,
136 +)
137 +class DiffusionGemmaForConditionalGeneration(
138 + nn.Module,
139 + SupportsMultiModal,
140 + SupportsQuant,
141 + SupportsPP,
142 +):
143 + """DiffusionGemma for vLLM.
144 +
145 + Single Gemma4 backbone that switches between encoder and decoder mode.
146 + The encoder path uses standard Gemma4 layers (causal attention, KV write).
147 + The decoder path uses the same weights with bidirectional attention and
148 + KV read-only, plus self-conditioning.
149 +
150 + Always includes a vision tower (same as Gemma4) for image understanding.
151 +
152 + In practice, the model's forward() dispatches based on the `mode` kwarg
153 + set by DiffusionGemmaModelState.prepare_inputs().
154 + """
155 +
156 + hf_to_vllm_mapper = WeightsMapper(
157 + orig_to_new_prefix={
158 + "model.decoder.": "model.",
159 + "model.encoder.language_model.": "model.",
160 + "model.encoder.vision_tower.": "vision_tower.",
161 + "model.encoder.embed_vision.": "embed_vision.",
162 + },
163 + orig_to_new_substr={
164 + ".experts.": ".moe.experts.",
165 + },
166 + )
167 +
168 + packed_modules_mapping = {
169 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
170 + "gate_up_proj": ["gate_proj", "up_proj"],
171 + }
172 +
173 + @staticmethod
174 + def get_model_state_cls():
175 + return DiffusionGemmaModelState
176 +
177 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
178 + super().__init__()
179 + config = vllm_config.model_config.hf_config
180 + text_config = vllm_config.model_config.hf_text_config
181 + self.config = config
182 + self.model_dtype = vllm_config.model_config.dtype
183 +
184 + # DiffusionGemma's full-attention layers have NO v_proj — V is
185 + # computed from k_proj's output (`value_states = key_states` before
186 + # k_norm in `DiffusionGemmaDecoderTextAttention.forward`). This is
187 + # the "k_eq_v" variant in our Gemma4 backbone. The checkpoint has no
188 + # v_proj weights for full-attention layers; without this flag they
189 + # would silently load with random V projections.
190 + text_config.attention_k_eq_v = True
191 +
192 + # ---- Vision tower ----
193 + vision_config = getattr(config, "vision_config", None)
194 + if vision_config is not None:
195 + quant_config = vllm_config.quant_config
196 + if quant_config and quant_config.get_name() in [
197 + "bitsandbytes",
198 + "torchao",
199 + "compressed-tensors",
200 + ]:
201 + tower_quant = quant_config
202 + else:
203 + quantizable = (
204 + vision_config.hidden_size % 64 == 0
205 + and vision_config.intermediate_size % 64 == 0
206 + )
207 + tower_quant = quant_config if quantizable else None
208 +
209 + with self._mark_tower_model(vllm_config, {"image", "video"}):
210 + self.vision_tower = AutoModel.from_config(config=vision_config)
211 + self.embed_vision = Gemma4MultimodalEmbedder(
212 + vision_config,
213 + text_config,
214 + quant_config=tower_quant,
215 + prefix=maybe_prefix(prefix, "embed_vision"),
216 + )
217 + recursive_replace_linear(
218 + self.vision_tower,
219 + tower_quant,
220 + prefix=maybe_prefix(prefix, "vision_tower"),
221 + )
222 + else:
223 + self.vision_tower = None
224 + self.embed_vision = None
225 +
226 + # ---- Language backbone (Gemma4Model) ----
227 + # Use maybe_prefix to ensure correct weight name prefixes for
228 + # quantization. The quantization config uses hf_to_vllm_mapper to
229 + # match checkpoint weight names to model parameter names.
230 + self.model = Gemma4Model(
231 + vllm_config=vllm_config,
232 + prefix=maybe_prefix(prefix, "model"),
233 + )
234 +
235 + self.lm_head = ParallelLMHead(
236 + num_embeddings=text_config.vocab_size,
237 + embedding_dim=text_config.hidden_size,
238 + )
239 +
240 + if text_config.tie_word_embeddings:
241 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
242 +
243 + # HF DiffusionGemma applies the final-logit softcap in fp32, before
244 + # any other processing. Do it manually in `compute_logits` so the
245 + # LogitsProcessor only handles the lm_head GEMM.
246 + self.final_logit_softcapping = getattr(
247 + text_config, "final_logit_softcapping", None
248 + )
249 + self.logits_processor = LogitsProcessor(
250 + text_config.vocab_size,
251 + soft_cap=None,
252 + )
253 +
254 + sc_size = (
255 + getattr(config, "self_conditioning_size", None)
256 + or text_config.intermediate_size
257 + )
258 + self.self_conditioning = DiffusionGemmaSelfConditioning(
259 + hidden_size=text_config.hidden_size,
260 + self_conditioning_size=sc_size,
261 + eps=getattr(text_config, "rms_norm_eps", 1e-6),
262 + )
263 +
264 + self.make_empty_intermediate_tensors = (
265 + self.model.make_empty_intermediate_tensors
266 + )
267 +
268 + def compute_self_conditioning(
269 + self,
270 + inputs_embeds: torch.Tensor,
271 + probs: torch.Tensor,
272 + ) -> torch.Tensor:
273 + embed_weight = self.model.embed_tokens.weight
274 + soft_embeds = torch.matmul(
275 + probs.to(embed_weight.dtype), embed_weight
276 + ) * self.model.normalizer.to(inputs_embeds.dtype)
277 + return self.self_conditioning(inputs_embeds, soft_embeds)
278 +
279 + # ------------------------------------------------------------------ #
280 + # Multimodal: reuse Gemma4's image parsing, processing & embedding
281 + # ------------------------------------------------------------------ #
282 + # The vision tower, pooler, embed_vision, and their processing logic
283 + # are architecturally identical to Gemma4. Delegate to avoid
284 + # maintaining a duplicate copy.
285 +
286 + _parse_and_validate_image_input = (
287 + Gemma4ForConditionalGeneration._parse_and_validate_image_input
288 + )
289 + _parse_and_validate_video_input = (
290 + Gemma4ForConditionalGeneration._parse_and_validate_video_input
291 + )
292 + _parse_and_validate_multimodal_inputs = (
293 + Gemma4ForConditionalGeneration._parse_and_validate_multimodal_inputs
294 + )
295 + _encoder_chunk = staticmethod(Gemma4ForConditionalGeneration._encoder_chunk)
296 + _process_image_input = Gemma4ForConditionalGeneration._process_image_input
297 + _process_video_input = Gemma4ForConditionalGeneration._process_video_input
298 + embed_multimodal = Gemma4ForConditionalGeneration.embed_multimodal
299 +
300 + def get_mm_mapping(self) -> MultiModelKeys:
301 + """Get the module prefix mapping for multimodal models."""
302 + return MultiModelKeys.from_string_field(
303 + language_model="model",
304 + connector=["embed_vision"],
305 + tower_model=["vision_tower"],
306 + )
307 +
308 + # ------------------------------------------------------------------ #
309 + # Forward
310 + # ------------------------------------------------------------------ #
311 +
312 + def forward(
313 + self,
314 + input_ids: torch.Tensor,
315 + positions: torch.Tensor,
316 + intermediate_tensors: Any | None = None,
317 + inputs_embeds: torch.Tensor | None = None,
318 + **kwargs: Any,
319 + ) -> torch.Tensor:
320 + if intermediate_tensors is not None:
321 + inputs_embeds = None
322 + return self.model(
323 + input_ids=input_ids,
324 + positions=positions,
325 + intermediate_tensors=intermediate_tensors,
326 + inputs_embeds=inputs_embeds,
327 + **kwargs,
328 + )
329 +
330 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
331 + logits = self.logits_processor(self.lm_head, hidden_states)
332 + if logits is not None and self.final_logit_softcapping is not None:
333 + logits = _softcap_logits(logits, self.final_logit_softcapping)
334 + return logits
335 +
336 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
337 + """Load weights from checkpoint.
338 +
339 + Checkpoint layout (HF DiffusionGemma):
340 + model.encoder.vision_tower.* → vision tower
341 + model.encoder.embed_vision.* → vision embedder
342 + model.encoder.language_model.layers.* → backbone
343 + model.decoder.layers.* → backbone (tied)
344 + model.decoder.embed_tokens.* → embeddings
345 + model.decoder.self_conditioning.* → self-conditioning MLP
346 + lm_head.* → LM head (tied)
347 +
348 + We load encoder weights into our single ``Gemma4Model`` backbone,
349 + skip duplicate decoder backbone weights, handle vision tower and
350 + self-conditioning separately.
351 + """
352 +
353 + sc_params = dict(
354 + (n, p)
355 + for n, p in self.named_parameters()
356 + if n.startswith("self_conditioning.")
357 + )
358 +
359 + # Collect vision tower + embedder parameters AND buffers for manual
360 + # loading. The HF vision tower registers std_bias / std_scale as
361 + # buffers (not parameters) when config.standardize is True, so we
362 + # must include named_buffers() to avoid "not found in model" warnings.
363 + vision_params: dict[str, torch.Tensor] = {}
364 + for n, p in self.named_parameters():
365 + if n.startswith(("vision_tower.", "embed_vision.")):
366 + vision_params[n] = p
367 + for n, b in self.named_buffers():
368 + if n.startswith(("vision_tower.", "embed_vision.")):
369 + vision_params[n] = b
370 +
371 + def _remap_weights():
372 + # Use full weight names (including suffixes like .weight_scale,
373 + # .weight_packed) for dedup instead of just the base layer name. Critical
374 + # for quantized checkpoints where each weight has multiple tensors;
375 + # tracking only base names skips scales as duplicates.
376 + seen_weights: set[str] = set()
377 + for name, weight in weights:
378 + # Self-conditioning lives under model.decoder.self_conditioning.*
379 + # in the checkpoint but at self_conditioning.* in our model.
380 + if "self_conditioning" in name:
381 + sc_name = name.split("self_conditioning.", 1)[1]
382 + sc_name = "self_conditioning." + sc_name
383 + if sc_name in sc_params:
384 + sc_params[sc_name].data.copy_(weight)
385 + continue
386 +
387 + # Vision tower: model.encoder.vision_tower.* → vision_tower.*
388 + # In HF, the vision tower is a sibling of language_model
389 + # under the encoder module.
390 + if name.startswith("model.encoder.vision_tower."):
391 + vt_name = name[len("model.encoder.") :]
392 + if vt_name in vision_params:
393 + vision_params[vt_name].data.copy_(weight)
394 + else:
395 + logger.warning(
396 + "Vision tower weight %s (mapped to %s) not found in model",
397 + name,
398 + vt_name,
399 + )
400 + continue
401 +
402 + # Vision embedder: model.encoder.embed_vision.* → embed_vision.*
403 + if name.startswith("model.encoder.embed_vision."):
404 + ev_name = name[len("model.encoder.") :]
405 + if ev_name in vision_params:
406 + vision_params[ev_name].data.copy_(weight)
407 + else:
408 + logger.warning(
409 + "Embed vision weight %s (mapped to %s) not found in model",
410 + name,
411 + ev_name,
412 + )
413 + continue
414 +
415 + # Skip vestigial embed_vision.embedding weights.
416 + if "embed_vision.embedding." in name:
417 + continue
418 +
419 + # Encoder backbone → model.*
420 + if name.startswith("model.encoder.language_model."):
421 + name = name.replace("model.encoder.language_model.", "model.")
422 + # Decoder backbone → model.* (skip exact duplicates)
423 + elif name.startswith("model.decoder."):
424 + name = name.replace("model.decoder.", "model.")
425 +
426 + # Skip only if we've seen the exact same weight name (including scales)
427 + if name in seen_weights:
428 + continue
429 + seen_weights.add(name)
430 + yield name, weight
431 +
432 + # Delegate to Gemma4ForCausalLM.load_weights for the backbone,
433 + # which handles stacked params, MoE, k_eq_v, etc.
434 + # Temporarily set self.config to text_config since Gemma4's
435 + # load_weights expects it (e.g. tie_word_embeddings, layer_types).
436 + from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
437 +
438 + saved_config = self.config
439 + self.config = self.model.config
440 + try:
441 + Gemma4ForCausalLM.load_weights(self, _remap_weights())
442 + finally:
443 + self.config = saved_config
444 +
445 + @classmethod
446 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
447 + if modality == "image":
448 + return "<image_soft_token>"
449 + if modality == "video":
450 + return "<|video|>"
451 + raise ValueError(f"Unsupported modality: {modality}")
452 +
453 +
454 +@torch.compile(dynamic=True)
455 +def _compute_num_rejected(
456 + num_logits: torch.Tensor,
457 + num_sampled: torch.Tensor,
458 + query_start_loc: torch.Tensor,
459 +) -> torch.Tensor:
460 + query_lens = query_start_loc[1:] - query_start_loc[:-1]
461 + num_rejected = num_logits - num_sampled
462 + is_denoise = (num_logits > 0) & (num_sampled == 0)
463 + return torch.where(is_denoise, query_lens, num_rejected)
464 +
465 +
466 +@torch.compile(dynamic=True)
467 +def _compiled_sample_step(
468 + # Logits from the model [num_decode * CL, vocab]
469 + logits: torch.Tensor,
470 + # Request mapping
471 + decode_slots: torch.Tensor, # [num_decode] int64 → slot indices
472 + decode_idx: torch.Tensor, # [num_decode] int64 → position in num_reqs
473 + all_slots: torch.Tensor, # [num_reqs] int64 → all slot indices
474 + valid_canvas_len: torch.Tensor, # [num_decode] int64 → real canvas length (<=CL)
475 + # State tensors (modified in-place)
476 + canvas: torch.Tensor, # [max_num_reqs, CL]
477 + argmax_canvas: torch.Tensor, # [max_num_reqs, CL]
478 + step_tensor: torch.Tensor, # [max_num_reqs]
479 + is_encoder_phase: torch.Tensor, # [max_num_reqs]
480 + confident_tensor: torch.Tensor, # [max_num_reqs]
481 + sc_embeds: torch.Tensor, # [max_num_reqs, CL, hidden]
482 + embed_weight: torch.Tensor, # [vocab, hidden]
483 + normalizer: torch.Tensor,
484 + history: torch.Tensor, # [max_num_reqs, ST, CL]
485 + history_len_tensor: torch.Tensor, # [max_num_reqs]
486 + # Output tensors (modified in-place)
487 + sampled: torch.Tensor, # [num_reqs, CL]
488 + num_sampled: torch.Tensor, # [num_reqs]
489 + draft_tokens: torch.Tensor, # [max_num_reqs, >=CL]
490 + # Scalar config
491 + max_denoising_steps: float,
492 + t_min: float,
493 + t_max: float,
494 + confidence_threshold: float,
495 + vocab_size: int,
496 + CL: int,
497 + ST: int,
498 + # Sampler config
499 + entropy_bound: float,
500 +) -> torch.Tensor:
501 + """Compiled decode step: temperature → Gumbel sample → probs/confidence →
502 + accept/renoise → convergence, all as vectorized PyTorch ops.
503 +
504 + Returns the temperature-scaled logits ``[num_decode, CL, vocab]`` so the
505 + caller can compute logprobs outside the compiled region."""
506 + num_decode = decode_slots.shape[0]
507 + device = decode_slots.device
508 +
509 + # Clear outputs so prefill / non-decode slots report 0 (decode slots are
510 + # overwritten below).
511 + sampled.zero_()
512 + num_sampled.zero_()
513 +
514 + # ---- Phase 1: Temperature schedule ----
515 + steps_f = step_tensor[decode_slots].float()
516 + remaining = (max_denoising_steps - steps_f).clamp(min=1.0)
517 + temp = t_min + (t_max - t_min) * (remaining / max_denoising_steps)
518 +
519 + # ---- Phase 2: Temperature scaling + Gumbel-max sampling ----
520 + logits_3d = logits.reshape(num_decode, CL, -1).float()
521 + scaled = logits_3d / temp[:, None, None].clamp(min=1e-10)
522 +
523 + # Gumbel-max trick: argmax(logits/T + Gumbel) ~ sample from softmax(logits/T)
524 + u = torch.rand_like(scaled).clamp(min=1e-20)
525 + gumbel = -torch.log(-torch.log(u))
526 + # Zero noise when temp==0 (greedy)
527 + noisy = scaled + gumbel * (temp[:, None, None] > 0).float()
528 + new_tokens = noisy.view(-1, noisy.shape[-1]).argmax(dim=-1).view(num_decode, CL)
529 + argmax_tokens = (
530 + scaled.view(-1, scaled.shape[-1]).argmax(dim=-1).view(num_decode, CL)
531 + )
532 +
533 + # ---- Phase 3: Probs, self-conditioning, confidence ----
534 + log_probs = scaled.log_softmax(dim=-1)
535 + probs = log_probs.exp()
536 +
537 + token_entropy = -(probs * log_probs).sum(dim=-1) # [num_decode, CL]
538 + # A canvas truncated near max_model_len is zero-padded up to CL by the
539 + # caller; those padded rows are uniform (max entropy, argmax 0), so they
540 + # never trigger early convergence and are stable, and only the real
541 + # ``valid_canvas_len`` tokens are committed (num_sampled below).
542 + mean_entropy = token_entropy.mean(dim=-1) # [num_decode]
543 + confident_tensor[decode_slots] = mean_entropy < confidence_threshold
544 +
545 + # ---- Phase 4: Entropy-bound acceptance mask ----
546 + sorted_ent, sorted_idx = torch.sort(token_entropy, dim=-1)
547 + cumsum_ent = torch.cumsum(sorted_ent, dim=-1)
548 + cummax_ent = torch.cummax(sorted_ent, dim=-1).values
549 + sorted_mask = (cumsum_ent - cummax_ent) <= entropy_bound
550 + eb_mask = torch.zeros_like(sorted_mask)
551 + eb_mask.scatter_(1, sorted_idx, sorted_mask)
552 +
553 + # ---- Phase 5: Post-sample ----
554 + is_commit = is_encoder_phase[decode_slots] # [num_decode]
555 + is_denoise = ~is_commit
556 + cur_step = step_tensor[decode_slots].float()
557 +
558 + # Step update: +1 for denoise, reset to 0 for commit
559 + new_step_val = torch.where(
560 + is_denoise,
561 + (cur_step + 1).to(step_tensor.dtype),
562 + step_tensor.new_zeros(num_decode),
563 + )
564 + step_tensor[decode_slots] = new_step_val
565 +
566 + # Random tokens for renoise / canvas reinit
567 + random_tokens = torch.randint(
568 + 0, vocab_size, (num_decode, CL), device=device, dtype=canvas.dtype
569 + )
570 +
571 + # Compute denoise canvas (accept/renoise)
572 + denoise_canvas = torch.where(eb_mask, new_tokens, random_tokens)
573 +
574 + # Canvas: commit → random reinit, denoise → accept/renoise result
575 + canvas[decode_slots] = torch.where(
576 + is_commit.unsqueeze(1), random_tokens, denoise_canvas
577 + )
578 +
579 + # History: write argmax_tokens for denoise requests at circular position
580 + hist_len = history_len_tensor[decode_slots]
581 + write_pos = hist_len % ST
582 + for i in range(ST):
583 + write_here = ((write_pos == i) & is_denoise).unsqueeze(1)
584 + history[decode_slots, i] = torch.where(
585 + write_here, argmax_tokens, history[decode_slots, i]
586 + )
587 +
588 + # Argmax canvas: update for denoise, preserve for commit
589 + argmax_canvas[decode_slots] = torch.where(
590 + is_denoise.unsqueeze(1), argmax_tokens, argmax_canvas[decode_slots]
591 + )
592 +
593 + # History length: increment for denoise, reset for commit
594 + new_hist_len = torch.where(is_denoise, hist_len + 1, hist_len.new_zeros(num_decode))
595 + history_len_tensor[decode_slots] = new_hist_len
596 +
597 + # Sampled output: commit → emit argmax_canvas, denoise → 0 (pre-zeroed)
598 + sampled[decode_idx] = argmax_canvas[decode_slots].to(
599 + sampled.dtype
600 + ) * is_commit.unsqueeze(1).to(sampled.dtype)
601 + # Commit only the real canvas length (== CL except for a canvas truncated
602 + # near max_model_len); the padded tail positions are never emitted.
603 + num_sampled[decode_idx] = is_commit.to(num_sampled.dtype) * valid_canvas_len.to(
604 + num_sampled.dtype
605 + )
606 +
607 + # ---- Phase 6: Stability + convergence ----
608 + ref = history[decode_slots, 0]
609 + mismatch = torch.zeros(num_decode, device=device, dtype=torch.int32)
610 + for h in range(1, ST):
611 + mismatch = mismatch + (ref != history[decode_slots, h]).sum(dim=-1).int()
612 + stable = mismatch == 0
613 +
614 + step_after = step_tensor[decode_slots]
615 + converged = (stable & confident_tensor[decode_slots] & (new_hist_len >= ST)) | (
616 + step_after >= max_denoising_steps
617 + )
618 + # Commit done → denoise next (False); denoise converged → commit next (True)
619 + is_encoder_phase[decode_slots] = torch.where(
620 + is_commit, is_commit.new_zeros(num_decode), converged
621 + )
622 +
623 + # SC soft embedding: store ``probs @ embed_weight`` (the value the next step's
624 + # self-conditioning MLP consumes) only for slots that will denoise next — i.e.
625 + # this step denoised AND it isn't about to commit (is_encoder_phase now False).
626 + # Masking here (rather than in the consumer) lets _apply_self_conditioning read
627 + # sc_embeds directly. Storing the [.., hidden] soft embed instead of the full
628 + # [.., vocab] probs avoids a giant persistent buffer.
629 + sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None]
630 + soft_embeds = torch.matmul(probs.to(embed_weight.dtype), embed_weight) * normalizer
631 + sc_embeds[decode_slots] = soft_embeds * sc_keep
632 +
633 + # Overwrite canvas with argmax for newly converged denoise requests
634 + newly_converged = (converged & is_denoise).unsqueeze(1)
635 + canvas[decode_slots] = torch.where(
636 + newly_converged, argmax_canvas[decode_slots], canvas[decode_slots]
637 + )
638 +
639 + # ---- Phase 7: Copy canvas → draft_tokens for all slots ----
640 + draft_tokens[all_slots, :CL] = canvas[all_slots]
641 +
642 + return scaled
643 +
644 +
645 +class DiffusionGemmaRequestStates:
646 + """Pre-allocated GPU tensors for DiffusionGemma per-request state.
647 +
648 + Follows the indexed-slot pattern used by ``RequestState``.
649 + """
650 +
651 + def __init__(
652 + self,
653 + max_num_reqs: int,
654 + canvas_length: int,
655 + vocab_size: int,
656 + max_denoising_steps: int,
657 + device: torch.device,
658 + hidden_size: int,
659 + stability_threshold: int,
660 + ):
661 + self.max_num_reqs = max_num_reqs
662 + self.canvas_length = canvas_length
663 + self.vocab_size = vocab_size
664 + self.max_denoising_steps = max_denoising_steps
665 + self.stability_threshold = stability_threshold
666 + self.device = device
667 +
668 + self.is_encoder_phase = torch.zeros(
669 + max_num_reqs, dtype=torch.bool, device=device
670 + )
671 + # Canvas tokens [max_num_reqs, canvas_length]
672 + self.canvas = torch.zeros(
673 + max_num_reqs, canvas_length, dtype=torch.int64, device=device
674 + )
675 + # Step counter (counts up from 0 to max_denoising_steps)
676 + self.step = torch.zeros(
677 + max_num_reqs,
678 + dtype=torch.int32,
679 + device=device,
680 + )
681 + # Accepted canvas history for stability check
682 + self.accepted_canvas_history = torch.zeros(
683 + max_num_reqs,
684 + stability_threshold,
685 + canvas_length,
686 + dtype=torch.int64,
687 + device=device,
688 + )
689 + self.accepted_canvas_history_len = torch.zeros(
690 + max_num_reqs, dtype=torch.int32, device=device
691 + )
692 + # Latest argmax(processed_logits) per slot — what we COMMIT.
693 + # NOT `current_canvas` (which is the post-renoise stochastic input for
694 + # the next denoise step). We keep this separate from `canvas` because
695 + # canvas gets renoised in-place during denoise, while argmax_canvas is
696 + # the deterministic best-guess we ultimately emit.
697 + self.argmax_canvas = torch.zeros(
698 + max_num_reqs, canvas_length, dtype=torch.int64, device=device
699 + )
700 +
701 + # Per-slot prompt length (set by add_request).
702 + self.prompt_len = torch.zeros(
703 + max_num_reqs,
704 + dtype=torch.int32,
705 + device=device,
706 + )
707 +
708 + # Per-slot confidence flag, set by the sampler each step.
709 + self.confident = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
710 +
711 + # Per-slot self-conditioning soft embedding (probs @ embed_weight) from
712 + # the previous denoise step. Storing the [.., hidden] soft embed instead
713 + # of the full [.., vocab] distribution shrinks this buffer by
714 + # vocab/hidden (~170x) and moves the matmul to denoise time; the result
715 + # is identical (SC consumes probs @ embed_weight anyway).
716 + self.self_conditioning_embeds = torch.zeros(
717 + max_num_reqs, canvas_length, hidden_size, dtype=torch.float32, device=device
718 + )
719 +
720 + def init_canvas(self, slot_indices_np: np.ndarray) -> None:
721 + """Initialize canvas with random tokens for the given slots."""
722 + n = slot_indices_np.shape[0]
723 + self.canvas[slot_indices_np] = torch.randint(
724 + 0,
725 + self.vocab_size,
726 + (n, self.canvas_length),
727 + dtype=torch.int64,
728 + device=self.device,
729 + )
730 +
731 + def add_request(self, slot_idx: int) -> None:
732 + self.is_encoder_phase[slot_idx] = True
733 + self.init_canvas(torch.tensor([slot_idx], device=self.device))
734 + self.step[slot_idx] = 0
735 + self.accepted_canvas_history_len[slot_idx] = 0
736 + self.self_conditioning_embeds[slot_idx] = 0
737 +
738 + def remove_request(self, slot_idx: int) -> None:
739 + self.is_encoder_phase[slot_idx] = False
740 + self.accepted_canvas_history_len[slot_idx] = 0
741 + self.self_conditioning_embeds[slot_idx] = 0
742 +
743 +
744 +class DiffusionGemmaModelState(ModelState):
745 + """ModelState for DiffusionGemma.
746 +
747 + Single Gemma4 backbone in two modes:
748 + - encoder mode (num_draft_tokens == 0): causal attention, writes KV
749 + - decoder mode (num_draft_tokens > 0): bidirectional attention, reads KV
750 + """
751 +
752 + def __init__(
753 + self,
754 + vllm_config: VllmConfig,
755 + model: nn.Module,
756 + encoder_cache: Any,
757 + device: torch.device,
758 + ) -> None:
759 + self.vllm_config = vllm_config
760 + self.model_config = vllm_config.model_config
761 + self.scheduler_config = vllm_config.scheduler_config
762 + self.model = model
763 + self.device = device
764 +
765 + self.supports_mm_inputs = encoder_cache is not None
766 + self.max_num_reqs = self.scheduler_config.max_num_seqs
767 + self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
768 + self.max_model_len = self.model_config.max_model_len
769 + self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
770 + self.dtype = self.model_config.dtype
771 +
772 + if self.supports_mm_inputs:
773 + from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
774 + from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
775 +
776 + assert isinstance(encoder_cache, EncoderCache)
777 + self.encoder_cache = encoder_cache
778 + self.encoder_runner = EncoderRunner(
779 + model=self.model,
780 + max_num_tokens=self.max_num_tokens,
781 + hidden_size=self.inputs_embeds_size,
782 + encoder_cache=encoder_cache,
783 + dtype=self.dtype,
784 + device=self.device,
785 + )
786 +
787 + # Per-step MM data produced by get_mm_embeddings and consumed by
788 + # prepare_inputs. Stored as raw (mm_embeds, is_mm_embed) so that
789 + # prepare_inputs can call embed_input_ids directly into the
790 + # persistent _inputs_embeds_buf, avoiding the intermediate copy
791 + # through encoder_runner.inputs_embeds.
792 + self._pending_mm_embeds: tuple[list[torch.Tensor], torch.Tensor] | None = None
793 +
794 + diffusion_config = vllm_config.diffusion_config
795 + canvas_length = diffusion_config.canvas_length if diffusion_config else 32
796 +
797 + text_config = self.model_config.hf_text_config
798 + self.gen_config = self.model_config.try_get_generation_config()
799 + max_denoising_steps = (
800 + diffusion_config.max_denoising_steps if diffusion_config else None
801 + ) or self.gen_config.get("max_denoising_steps", 48)
802 + self.diffusion_states = DiffusionGemmaRequestStates(
803 + max_num_reqs=self.max_num_reqs,
804 + canvas_length=canvas_length,
805 + vocab_size=self.model_config.get_vocab_size(),
806 + max_denoising_steps=max_denoising_steps,
807 + device=device,
808 + hidden_size=text_config.hidden_size,
809 + stability_threshold=self.gen_config["stability_threshold"],
810 + )
811 + self._req_id_to_index: dict[str, int] = {}
812 +
813 + # Persistent buffer for per-request causal flags, updated in-place
814 + # so FULL CUDA graph replay sees the latest values.
815 + self._causal_buf = torch.zeros(
816 + self.max_num_reqs, dtype=torch.bool, device=device
817 + )
818 +
819 + # Persistent inputs_embeds buffer — required so FULL CUDA graph
820 + # capture and runtime point at the SAME memory address.
821 + # `prepare_dummy_inputs` (capture path) and `prepare_inputs` (runtime
822 + # path) both must hand the captured graph a tensor at this address.
823 + self._inputs_embeds_buf = torch.zeros(
824 + self.max_num_tokens,
825 + text_config.hidden_size,
826 + dtype=self.model_config.dtype,
827 + device=device,
828 + )
829 +
830 + def get_supported_generation_tasks(self):
831 + return ("generate",)
832 +
833 + def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None:
834 + diffusion_config = self.vllm_config.diffusion_config
835 + gen = self.gen_config
836 + sampler_cfg = gen.get("sampler_config") or {}
837 + if "EntropyBound" not in sampler_cfg.get("_cls_name", ""):
838 + raise ValueError("DiffusionGemma requires an EntropyBound sampler_config")
839 + entropy_bound = sampler_cfg.get("entropy_bound")
840 + if entropy_bound is None or entropy_bound <= 0:
841 + raise ValueError(
842 + f"entropy_bound must be a positive float (got {entropy_bound})"
843 + )
844 + return DiffusionSampler(
845 + sampler=sampler,
846 + diffusion_config=diffusion_config,
847 + vocab_size=self.model_config.get_vocab_size(),
848 + diffusion_states=self.diffusion_states,
849 + t_min=gen["t_min"],
850 + t_max=gen["t_max"],
851 + entropy_bound=entropy_bound,
852 + confidence_threshold=gen["confidence_threshold"],
853 + embed_weight=self.model.model.embed_tokens.weight,
854 + normalizer=self.model.model.normalizer,
855 + ), None
856 +
857 + def apply_staged_writes(self) -> None:
858 + pass
859 +
860 + def add_request(self, req_index: int, new_req_data: Any) -> None:
861 + self._req_id_to_index[new_req_data.req_id] = req_index
862 + self.diffusion_states.add_request(req_index)
863 + if not new_req_data.req_id.startswith("_warmup_"):
864 + prompt_len = len(new_req_data.prompt_token_ids)
865 + self.diffusion_states.prompt_len[req_index] = prompt_len
866 +
867 + def remove_request(self, req_id: str) -> None:
868 + idx = self._req_id_to_index.pop(req_id, None)
869 + if idx is not None:
870 + self.diffusion_states.remove_request(idx)
871 +
872 + def get_mm_embeddings(self, scheduled_encoder_inputs, input_batch):
873 + if not self.supports_mm_inputs:
874 + return None
875 +
876 + mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
877 + scheduled_encoder_inputs
878 + )
879 + if mm_kwargs:
880 + encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
881 + self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))
882 +
883 + mm_embeds, is_mm_embed = self.encoder_runner.gather_mm_embeddings(
884 + input_batch.req_ids,
885 + input_batch.num_tokens,
886 + input_batch.num_scheduled_tokens,
887 + input_batch.query_start_loc_np,
888 + input_batch.prefill_len_np,
889 + input_batch.num_computed_prefill_tokens_np,
890 + )
891 +
892 + if not mm_embeds:
893 + # No MM tokens in this batch (e.g. all-decode step).
894 + # prepare_inputs will use embed_input_ids (text-only) directly.
895 + self._pending_mm_embeds = None
896 + return None
897 +
898 + # Stash raw MM ingredients for prepare_inputs to merge directly
899 + # into the persistent buffer, avoiding the intermediate copy
900 + # through encoder_runner.inputs_embeds.
901 + self._pending_mm_embeds = (mm_embeds, is_mm_embed)
902 + return None
903 +
904 + def _apply_self_conditioning(
905 + self,
906 + decode_slots_np: np.ndarray,
907 + decode_idx_np: np.ndarray,
908 + query_start_loc_np: np.ndarray,
909 + inputs_embeds: torch.Tensor,
910 + sc_embeds: torch.Tensor,
911 + ) -> None:
912 + # One self-conditioning MLP call per decode request, over that request's
913 + # query span [start, end) = its canvas. The span is the full canvas (CL)
914 + # or, for the final canvas truncated near max_model_len, fewer than CL
915 + # positions. sc_embeds already holds probs @ embed_weight from the prior
916 + # denoise step, masked to zero by the sampler for slots not denoising
917 + # this step; only the MLP runs here. CPU metadata -> no GPU syncs.
918 + for slot, idx in zip(decode_slots_np.tolist(), decode_idx_np.tolist()):
919 + start = int(query_start_loc_np[idx])
920 + end = int(query_start_loc_np[idx + 1])
921 + canvas = slice(start, end)
922 + soft = sc_embeds[slot, : end - start]
923 + inputs_embeds[canvas] = self.model.self_conditioning(
924 + inputs_embeds[canvas], soft.to(inputs_embeds.dtype)
925 + )
926 +
927 + def prepare_inputs(self, input_batch, req_states) -> dict[str, Any]:
928 + states = self.diffusion_states
929 + num_tokens = input_batch.num_tokens
930 + num_reqs = input_batch.num_reqs
931 +
932 + # Write into the PERSISTENT inputs_embeds buffer so FULL CUDA graph
933 + # replay sees the latest values at the captured address.
934 + num_tokens_padded = input_batch.num_tokens_after_padding
935 + inputs_embeds = self._inputs_embeds_buf[:num_tokens_padded]
936 +
937 + # Populate embeddings: merge MM features when available,
938 + # otherwise embed input_ids as text-only.
939 + input_ids = input_batch.input_ids[:num_tokens]
940 + if self._pending_mm_embeds is not None:
941 + mm_embeds, is_mm_embed = self._pending_mm_embeds
942 + self._pending_mm_embeds = None
943 + inputs_embeds[:num_tokens].copy_(
944 + self.model.embed_input_ids(
945 + input_ids,
946 + multimodal_embeddings=mm_embeds,
947 + is_multimodal=is_mm_embed,
948 + )
949 + )
950 + else:
951 + inputs_embeds[:num_tokens].copy_(self.model.embed_input_ids(input_ids))
952 +
953 + # Apply self-conditioning ONLY for denoising decode requests.
954 + if input_batch.num_draft_tokens > 0 and self._req_id_to_index:
955 + slots_np = input_batch.idx_mapping_np[:num_reqs]
956 + num_logits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1])
957 + is_decode_indices_np = np.where(num_logits_np > 0)[0]
958 + self._apply_self_conditioning(
959 + slots_np[is_decode_indices_np],
960 + is_decode_indices_np,
961 + input_batch.query_start_loc_np,
962 + inputs_embeds,
963 + states.self_conditioning_embeds,
964 + )
965 +
966 + return {"inputs_embeds": inputs_embeds}
967 +
968 + def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
969 + # CUDA graph capture path — return a slice of the SAME persistent
970 + # inputs_embeds buffer that `prepare_inputs` writes to at runtime,
971 + # so the captured graph and runtime point to identical addresses.
972 + return {"inputs_embeds": self._inputs_embeds_buf[:num_tokens]}
973 +
974 + def postprocess_state(self, idx_mapping, num_sampled) -> None:
975 + return None
976 +
977 + def prepare_attn(
978 + self,
979 + input_batch,
980 + cudagraph_mode,
981 + block_tables,
982 + slot_mappings,
983 + attn_groups,
984 + kv_cache_config,
985 + for_capture=False,
986 + ) -> dict[str, Any]:
987 + if cudagraph_mode == CUDAGraphMode.FULL:
988 + num_reqs = input_batch.num_reqs_after_padding
989 + num_tokens = input_batch.num_tokens_after_padding
990 + else:
991 + num_reqs = input_batch.num_reqs
992 + num_tokens = input_batch.num_tokens
993 +
994 + query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
995 + max_query_len = input_batch.num_scheduled_tokens.max().item()
996 +
997 + # Per-request causal mode: encoder (commit) = causal,
998 + # denoise = bidirectional. Pass GPU tensor so the attention
999 + # backend can handle mixed batches.
1000 + actual_num_reqs = input_batch.num_reqs
1001 + slots = input_batch.idx_mapping[:actual_num_reqs]
1002 + # Invariant: the sampler flips is_encoder_phase to False only after a
1003 + # request's FINAL prompt chunk, so a prompt spanning multiple chunks
1004 + # (longer than the token budget) stays causal for every chunk.
1005 + self._causal_buf[:actual_num_reqs] = self.diffusion_states.is_encoder_phase[
1006 + slots
1007 + ]
1008 + if actual_num_reqs < num_reqs:
1009 + self._causal_buf[actual_num_reqs:num_reqs] = False
1010 + causal: bool | torch.Tensor = self._causal_buf[:num_reqs]
1011 +
1012 + return build_attn_metadata(
1013 + attn_groups=attn_groups,
1014 + num_reqs=num_reqs,
1015 + num_tokens=num_tokens,
1016 + query_start_loc_gpu=input_batch.query_start_loc,
1017 + query_start_loc_cpu=query_start_loc_cpu,
1018 + max_query_len=max_query_len,
1019 + seq_lens=input_batch.seq_lens,
1020 + max_seq_len=self.max_model_len,
1021 + block_tables=block_tables,
1022 + slot_mappings=slot_mappings,
1023 + kv_cache_config=kv_cache_config,
1024 + causal=causal,
1025 + )
1026 +
1027 + num_new_sampled_tokens_per_step: int = 0
1028 +
1029 +
1030 +# Penalty stub for the diffusion path: the runner reads
1031 +# penalties_state.output_bin_counts, and post_update treats None as
1032 +# "no penalty bookkeeping".
1033 +_NO_PENALTIES_STATE = SimpleNamespace(output_bin_counts=None)
1034 +
1035 +
1036 +class DiffusionSampler:
1037 + """Batched accept/renoise sampler for DiffusionGemma.
1038 +
1039 + Follows the same structure as ``vllm.v1.worker.gpu.sample.sampler.Sampler``:
1040 + decomposed into named methods, all GPU state in pre-allocated buffers,
1041 + no GPU→CPU syncs on the hot path.
1042 + """
1043 +
1044 + def __init__(
1045 + self,
1046 + sampler: Any,
1047 + diffusion_config: Any,
1048 + vocab_size: int,
1049 + diffusion_states: DiffusionGemmaRequestStates | None = None,
1050 + *,
1051 + confidence_threshold: float,
1052 + t_min: float,
1053 + t_max: float,
1054 + entropy_bound: float,
1055 + embed_weight: torch.Tensor,
1056 + normalizer: torch.Tensor,
1057 + ):
1058 + self.sampling_states = sampler.sampling_states
1059 + self.req_states = sampler.req_states
1060 + # Self-conditioning soft embed = probs @ embed_weight * normalizer,
1061 + # computed in the sampler (see _compiled_sample_step).
1062 + self.embed_weight = embed_weight
1063 + self.normalizer = normalizer
1064 + self.canvas_length = (
1065 + diffusion_config.canvas_length if diffusion_config is not None else 32
1066 + )
1067 + self.t_min = t_min
1068 + self.t_max = t_max
1069 + self.confidence_threshold = confidence_threshold
1070 + self.vocab_size = vocab_size
1071 + self.diffusion_states = diffusion_states
1072 + self.entropy_bound = entropy_bound
1073 +
1074 + max_num_reqs = diffusion_states.max_num_reqs
1075 + device = diffusion_states.device
1076 + self._sampled = torch.zeros(
1077 + max_num_reqs,
1078 + self.canvas_length,
1079 + dtype=torch.int32,
1080 + device=device,
1081 + )
1082 + self._num_sampled = torch.zeros(
1083 + max_num_reqs,
1084 + dtype=torch.int32,
1085 + device=device,
1086 + )
1087 + self._decode_slots = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
1088 + self._decode_idx = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
1089 + self._query_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
1090 + self._num_logits = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
1091 +
1092 + # Per-slot stash for logprobs computed on the converging denoise step.
1093 + # Populated after the post-sample kernel detects convergence; consumed
1094 + # on the subsequent commit step when num_sampled=CANVAS_LEN.
1095 + self._pending_logprobs: dict[int, LogprobsTensors] = {}
1096 +
1097 + def add_request(self, req_idx: int, prompt_len: int, sampling_params: Any) -> None:
1098 + if use_penalty(sampling_params):
1099 + logger.warning_once(
1100 + "DiffusionGemma does not support repetition/frequency/presence "
1101 + "penalties; ignoring them for this request."
1102 + )
1103 + # Purge any stale logprobs stashed under this slot by a prior request
1104 + # that was aborted between its converging denoise and commit steps.
1105 + self._pending_logprobs.pop(req_idx, None)
1106 + self.sampling_states.add_request(req_idx, sampling_params)
1107 +
1108 + def apply_staged_writes(self) -> None:
1109 + self.sampling_states.apply_staged_writes()
1110 +
1111 + @property
1112 + def penalties_state(self):
1113 + # Diffusion applies no penalties. The runner reads
1114 + # penalties_state.output_bin_counts, so expose a stub holding None;
1115 + # post_update treats None bin counts as "no penalty bookkeeping".
1116 + return _NO_PENALTIES_STATE
1117 +
1118 + # ------------------------------------------------------------------
1119 + # Prefill
1120 + # ------------------------------------------------------------------
1121 +
1122 + def _finish_prefills(
1123 + self, input_batch: Any, prefill_indices_np: np.ndarray
1124 + ) -> None:
1125 + """Transition requests whose prompt completes this step to denoising.
1126 +
1127 + Initializes their canvas, seeds draft tokens, and flips
1128 + is_encoder_phase to False. Mid-chunk requests (prompt longer than the
1129 + token budget) are left untouched so is_encoder_phase stays True and
1130 + prepare_attn keeps causal attention for their remaining chunks.
1131 + """
1132 + states = self.diffusion_states
1133 + done_prefill_np = (
1134 + input_batch.num_computed_prefill_tokens_np[prefill_indices_np]
1135 + + input_batch.num_scheduled_tokens[prefill_indices_np]
1136 + >= input_batch.prefill_len_np[prefill_indices_np]
1137 + )
1138 + ps = input_batch.idx_mapping_np[prefill_indices_np[done_prefill_np]]
1139 + if len(ps) == 0:
1140 + return
1141 + states.init_canvas(ps)
1142 + self.req_states.draft_tokens[ps, : self.canvas_length] = states.canvas[ps]
1143 + ps_gpu = async_copy_to_gpu(
1144 + ps.astype(np.int64), device=states.is_encoder_phase.device
1145 + )
1146 + states.is_encoder_phase.index_fill_(0, ps_gpu, False)
1147 +
1148 + def _handle_prefill(
1149 + self,
1150 + input_batch: Any,
1151 + device: torch.device,
1152 + ) -> SamplerOutput:
1153 + num_reqs = input_batch.num_reqs
1154 + self._finish_prefills(input_batch, np.arange(num_reqs))
1155 + sampled = self._sampled[:num_reqs, :1]
1156 + sampled.zero_()
1157 + num_sampled = self._num_sampled[:num_reqs]
1158 + num_sampled.zero_()
1159 + return SamplerOutput(
1160 + sampled_token_ids=sampled,
1161 + logprobs_tensors=None,
1162 + num_nans=None,
1163 + num_sampled=num_sampled,
1164 + num_rejected=num_sampled,
1165 + )
1166 +
1167 + # ------------------------------------------------------------------
1168 + # Decode helpers
1169 + # ------------------------------------------------------------------
1170 +
1171 + def _build_output(
1172 + self,
1173 + input_batch: Any,
1174 + sampled: torch.Tensor,
1175 + num_sampled: torch.Tensor,
1176 + per_req_nlogits_np: np.ndarray,
1177 + device: torch.device,
1178 + logprobs_tensors: LogprobsTensors | None = None,
1179 + ) -> SamplerOutput:
1180 + """Compute num_rejected and build SamplerOutput."""
1181 + num_reqs = input_batch.num_reqs
1182 +
1183 + self._query_lens.np[:num_reqs] = np.diff(
1184 + input_batch.query_start_loc_np[: num_reqs + 1]
1185 + )
1186 + self._num_logits.np[:num_reqs] = per_req_nlogits_np
1187 + self._query_lens.copy_to_uva()
1188 + self._num_logits.copy_to_uva()
1189 +
1190 + num_rejected = _compute_num_rejected(
1191 + self._num_logits.gpu[:num_reqs],
1192 + num_sampled,
1193 + input_batch.query_start_loc[: num_reqs + 1],
1194 + )
1195 +
1196 + return SamplerOutput(
1197 + sampled_token_ids=sampled,
1198 + logprobs_tensors=logprobs_tensors,
1199 + num_nans=None,
1200 + num_sampled=num_sampled,
1201 + num_rejected=num_rejected,
1202 + )
1203 +
1204 + # ------------------------------------------------------------------
1205 + # Main entry point
1206 + # ------------------------------------------------------------------
1207 +
1208 + def __call__(
1209 + self,
1210 + logits: torch.Tensor,
1211 + input_batch: Any,
1212 + draft_logits: torch.Tensor | None = None,
1213 + ) -> SamplerOutput:
1214 + num_reqs = input_batch.num_reqs
1215 + device = logits.device
1216 +
1217 + if input_batch.num_draft_tokens == 0:
1218 + return self._handle_prefill(input_batch, device)
1219 +
1220 + # --- CPU/NumPy setup (outside compile): split decode vs prefill, init
1221 + # canvas for any new prefills, and stage decode slot indices to GPU. ---
1222 + states = self.diffusion_states
1223 + CL = self.canvas_length
1224 + slots_np = input_batch.idx_mapping_np[:num_reqs]
1225 + per_req_nlogits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1])
1226 +
1227 + decode_indices_np = np.where(per_req_nlogits_np > 0)[0]
1228 + prefill_indices_np = np.where(per_req_nlogits_np == 0)[0]
1229 + decode_slots_np = slots_np[decode_indices_np]
1230 +
1231 + if len(prefill_indices_np) > 0:
1232 + self._finish_prefills(input_batch, prefill_indices_np)
1233 +
1234 + num_decode = len(decode_indices_np)
1235 + self._decode_slots.np[:num_decode] = decode_slots_np
1236 + self._decode_idx.np[:num_decode] = decode_indices_np
1237 + self._decode_slots.copy_to_uva()
1238 + self._decode_idx.copy_to_uva()
1239 + decode_slots = self._decode_slots.gpu[:num_decode]
1240 + decode_idx = self._decode_idx.gpu[:num_decode]
1241 +
1242 + # Real canvas length per decode request. Equals CL except when a canvas
1243 + # was truncated near max_model_len, in which case the scheduler gave us
1244 + # fewer than CL logits for that request.
1245 + valid_canvas_len_np = per_req_nlogits_np[per_req_nlogits_np > 0]
1246 + valid_canvas_len = async_copy_to_gpu(
1247 + valid_canvas_len_np.astype(np.int64), device=device
1248 + )
1249 +
1250 + # Pad any truncated canvas back to CL so the uniform-CL sampler math
1251 + # holds. Phantom (padded) positions are zeroed → uniform logits → high
1252 + # entropy (no premature convergence) and argmax 0 (stable); they are
1253 + # never committed (num_sampled == real length).
1254 + if num_decode > 0 and valid_canvas_len_np.min() < CL:
1255 + ar = torch.arange(CL, device=device)
1256 + starts = valid_canvas_len.cumsum(0) - valid_canvas_len # row offset per req
1257 + valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1) # [num_decode, CL]
1258 + src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1)
1259 + logits = logits[src.reshape(-1)] * valid.reshape(-1, 1).to(logits.dtype)
1260 +
1261 + # Cleared inside _compiled_sample_step so prefill/non-decode slots stay 0.
1262 + sampled = self._sampled[:num_reqs]
1263 + num_sampled = self._num_sampled[:num_reqs]
1264 +
1265 + all_slots = input_batch.idx_mapping[:num_reqs]
1266 +
1267 + # Snapshot which slots are committing BEFORE the compiled step runs,
1268 + # since it mutates is_encoder_phase (commit→False, converge→True).
1269 + is_committing = states.is_encoder_phase[decode_slots].clone()
1270 +
1271 + # --- Single compiled call: temp → sample → probs → post-process ---
1272 + scaled = _compiled_sample_step(
1273 + logits,
1274 + decode_slots,
1275 + decode_idx,
1276 + all_slots,
1277 + valid_canvas_len,
1278 + # State
1279 + states.canvas,
1280 + states.argmax_canvas,
1281 + states.step,
1282 + states.is_encoder_phase,
1283 + states.confident,
1284 + states.self_conditioning_embeds,
1285 + self.embed_weight,
1286 + self.normalizer,
1287 + states.accepted_canvas_history,
1288 + states.accepted_canvas_history_len,
1289 + # Output
1290 + sampled,
1291 + num_sampled,
1292 + self.req_states.draft_tokens,
1293 + # Config
1294 + max_denoising_steps=float(states.max_denoising_steps),
1295 + t_min=self.t_min,
1296 + t_max=self.t_max,
1297 + confidence_threshold=self.confidence_threshold,
1298 + vocab_size=self.vocab_size,
1299 + CL=self.canvas_length,
1300 + ST=states.stability_threshold,
1301 + entropy_bound=self.entropy_bound,
1302 + )
1303 +
1304 + # --- Logprobs: stash on convergence, return on commit ---
1305 + slots_np = input_batch.idx_mapping_np[:num_reqs]
1306 + is_decode_np = per_req_nlogits_np > 0
1307 +
1308 + logprobs_tensors = None
1309 + max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np)
1310 + if max_num_logprobs >= 0:
1311 + # Denoise steps that just converged: the compiled step flipped
1312 + # is_encoder_phase from False→True. Detect as slots where
1313 + # is_encoder_phase is now True but is_committing was False.
1314 + converged_mask = states.is_encoder_phase[decode_slots]
1315 + just_converged = converged_mask & ~is_committing
1316 + if just_converged.any():
1317 + flat_logits = scaled.reshape(-1, scaled.shape[-1])
1318 + argmax_tokens = scaled.argmax(dim=-1)
1319 + for local_idx in just_converged.nonzero(as_tuple=True)[0]:
1320 + li = local_idx.item()
1321 + slot = decode_slots[local_idx]
1322 + # Stash only the real canvas positions (== CL unless this
1323 + # canvas was truncated near max_model_len); padded tail
1324 + # positions are never emitted.
1325 + k_i = int(valid_canvas_len_np[li])
1326 + start = li * CL
1327 + self._pending_logprobs[slot.item()] = compute_topk_logprobs(
1328 + flat_logits[start : start + k_i],
1329 + max_num_logprobs,
1330 + argmax_tokens[local_idx][:k_i],
1331 + )
1332 +
1333 + # Commit steps: is_committing was True at entry. Reassemble
1334 + # previously stashed logprobs and attach to SamplerOutput.
1335 + if is_committing.any() and self._pending_logprobs:
1336 + parts_ids, parts_lp, parts_ranks = [], [], []
1337 + cu_gen: list[int] = []
1338 + flat_offset = 0
1339 + for i in range(num_reqs):
1340 + cu_gen.append(flat_offset)
1341 + slot = int(slots_np[i])
1342 + if is_decode_np[i] and slot in self._pending_logprobs:
1343 + lp = self._pending_logprobs.pop(slot)
1344 + parts_ids.append(lp.logprob_token_ids)
1345 + parts_lp.append(lp.logprobs)
1346 + parts_ranks.append(lp.selected_token_ranks)
1347 + flat_offset += lp.logprobs.shape[0]
1348 + if parts_ids:
1349 + logprobs_tensors = LogprobsTensors(
1350 + logprob_token_ids=torch.cat(parts_ids),
1351 + logprobs=torch.cat(parts_lp),
1352 + selected_token_ranks=torch.cat(parts_ranks),
1353 + cu_num_generated_tokens=cu_gen,
1354 + )
1355 +
1356 + return self._build_output(
1357 + input_batch,
1358 + sampled,
1359 + num_sampled,
1360 + per_req_nlogits_np,
1361 + device,
1362 + logprobs_tensors=logprobs_tensors,
1363 + )
vllm/model_executor/models/gemma4.py
+1 -3
@@ -725,10 +725,8 @@ class Gemma4DecoderLayer(nn.Module):
725 if self.enable_moe_block:
726 hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states)
727
728 - # Router and MoE experts see the residual (pre-MLP state),
729 - # matching the HF transformers forward path
730 - router_logits = self.router(residual)
728 hidden_states_2 = self.pre_feedforward_layernorm_2(residual)
729 + router_logits = self.router(residual)
730 hidden_states_2 = self.moe(hidden_states_2, router_logits)
731 hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2)
732
vllm/model_executor/models/registry.py
+4
@@ -400,6 +400,10 @@ _MULTIMODAL_MODELS = {
400 "gemma3n_mm",
401 "Gemma3nForConditionalGeneration",
402 ),
403 + "DiffusionGemmaForBlockDiffusion": (
404 + "diffusion_gemma",
405 + "DiffusionGemmaForConditionalGeneration",
406 + ),
407 "Gemma4ForConditionalGeneration": ("gemma4_mm", "Gemma4ForConditionalGeneration"),
408 "Gemma4UnifiedForConditionalGeneration": (
409 "gemma4_unified",
vllm/tool_parsers/gemma4_tool_parser.py
+145 -39
@@ -20,9 +20,11 @@ import json
20 from collections.abc import Sequence
21
22 import regex as re
23 +from openai.types.responses import ToolChoiceFunction
24
25 from vllm.entrypoints.chat_utils import make_tool_call_id
26 from vllm.entrypoints.openai.chat_completion.protocol import (
27 + ChatCompletionNamedToolChoiceParam,
28 ChatCompletionRequest,
29 )
30 from vllm.entrypoints.openai.engine.protocol import (
@@ -343,6 +345,9 @@ class Gemma4ToolParser(ToolParser):
345 tool parsers.
346 """
347
348 + # Gemma4 emits native special-token tool calls, not generic JSON calls.
349 + supports_required_and_named = False
350 +
351 def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
352 super().__init__(tokenizer, tools)
353
@@ -390,6 +395,23 @@ class Gemma4ToolParser(ToolParser):
395 def adjust_request(
396 self, request: ChatCompletionRequest | ResponsesRequest
397 ) -> ChatCompletionRequest | ResponsesRequest:
398 + if request.tools:
399 + tc = request.tool_choice
400 + if tc == "required" or isinstance(
401 + tc,
402 + (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction),
403 + ):
404 + # Do NOT call super().adjust_request() for required/named tool
405 + # choice. The base implementation injects a JSON-array
406 + # `structured_outputs` schema and forces xgrammar guided
407 + # decoding, which conflicts with Gemma4's native
408 + # `<|tool_call>call:...` (non-JSON) tool syntax and crashes
409 + # EngineCore under MTP spec decode. The streaming/extraction
410 + # parser already handles the native output, so guided decoding
411 + # is skipped here (mirrors the GLM4 precedent).
412 + if request.tool_choice != "none":
413 + request.skip_special_tokens = False
414 + return request
415 request = super().adjust_request(request)
416 if request.tools and request.tool_choice != "none":
417 # Don't skip special tokens — <|tool_call> etc. are needed for
@@ -549,22 +571,40 @@ class Gemma4ToolParser(ToolParser):
571 return DeltaMessage(content=delta_text)
572 return None
573
552 - # Case 2: Starting a new tool call
553 - if start_count > prev_start_count and start_count > end_count:
554 - self.current_tool_id += 1
574 + # Case 2: One or more new tool calls started in this delta.
575 + # A single delta can batch several complete calls, so advance the
576 + # tool id once per newly-seen start token and allocate a tracking
577 + # slot for each.
578 + if start_count > prev_start_count:
579 + num_new = start_count - prev_start_count
580 + for _ in range(num_new):
581 + self.current_tool_id += 1
582 + self.streamed_args_for_tool.append("")
583 + self.prev_tool_call_arr.append({})
584 self.current_tool_name_sent = False
556 - self.streamed_args_for_tool.append("")
557 - self.prev_tool_call_arr.append({})
558 - logger.debug("Starting new tool call %d", self.current_tool_id)
559 - # Don't return yet — fall through to try parsing if there's
560 - # content after <|tool_call> in this same delta
561 - # (but usually it's just the token itself, so return None)
562 - if len(delta_text) <= len(self.tool_call_start_token):
585 + logger.debug(
586 + "Started %d new tool call(s); current_tool_id=%d",
587 + num_new,
588 + self.current_tool_id,
589 + )
590 + # Don't return yet if this delta also contains call payload or
591 + # the end marker; backends can batch one or more complete tool
592 + # calls into a single streaming chunk. Only wait for more text
593 + # when the delta is just the start token itself.
594 + if start_count > end_count and len(delta_text) <= len(
595 + self.tool_call_start_token
596 + ):
597 return None
598
565 - # Case 3: Tool call just ended
599 + # Case 3: One or more tool calls just ended (possibly several in a
600 + # single batched delta) — drain every newly-completed call.
601 if end_count > prev_end_count:
567 - return self._handle_tool_call_end(current_text)
602 + return self._handle_tool_call_end(
603 + current_text,
604 + prev_end_count=prev_end_count,
605 + end_count=end_count,
606 + start_count=start_count,
607 + )
608
609 # Case 4: In the middle of a tool call — parse partial content
610 if start_count > end_count:
@@ -652,45 +692,111 @@ class Gemma4ToolParser(ToolParser):
692
693 return None
694
655 - def _handle_tool_call_end(self, current_text: str) -> DeltaMessage | None:
656 - """Handle streaming when a tool call has just completed.
657 -
658 - Performs a final parse of the complete tool call and flushes
659 - any remaining un-streamed argument fragments.
695 + def _handle_tool_call_end(
696 + self,
697 + current_text: str,
698 + prev_end_count: int,
699 + end_count: int,
700 + start_count: int,
701 + ) -> DeltaMessage | None:
702 + """Handle streaming when one or more tool calls have just completed.
703 +
704 + A single streaming delta can batch several complete tool calls
705 + (``<|tool_call>...<tool_call|><|tool_call>...<tool_call|>``). Every
706 + call whose ``<tool_call|>`` end marker arrived in this delta — i.e.
707 + those with index in ``[prev_end_count, end_count)`` — is drained and
708 + emitted, with one ``DeltaToolCall`` per call in a single
709 + ``DeltaMessage`` (this matches the OpenAI streaming wire format, and
710 + the serving layer iterates over ``delta.tool_calls``).
711 +
712 + Per call:
713 +
714 + * If the function name was already streamed incrementally (the
715 + token-by-token path), only the remaining argument fragment is
716 + flushed as a diff.
717 + * If the call is seen complete for the first time in this delta (the
718 + batched-complete path), the id + name + full arguments JSON are
719 + emitted exactly once.
720 """
661 - if self.current_tool_id < 0 or self.current_tool_id >= len(
662 - self.prev_tool_call_arr
663 - ):
664 - logger.debug(
665 - "Tool call end detected but no active tool call (current_tool_id=%d)",
666 - self.current_tool_id,
667 - )
721 + # Parse the complete tool calls using regex for accuracy.
722 + all_matches = self.tool_call_regex.findall(current_text)
723 + if not all_matches:
724 + logger.debug("Tool call end detected but no complete tool call parsed yet.")
725 return None
726
670 - # Parse the complete tool call using regex for accuracy
671 - all_matches = self.tool_call_regex.findall(current_text)
672 - if self.current_tool_id < len(all_matches):
673 - _, args_str = all_matches[self.current_tool_id]
727 + deltas: list[DeltaToolCall] = []
728 + for idx in range(prev_end_count, end_count):
729 + if idx >= len(all_matches):
730 + break
731 + # Ensure the tracking arrays have a slot for this index (defensive;
732 + # Case 2 normally allocates these when the start token arrives).
733 + while len(self.prev_tool_call_arr) <= idx:
734 + self.prev_tool_call_arr.append({})
735 + self.streamed_args_for_tool.append("")
736 +
737 + func_name, args_str = all_matches[idx]
738 final_args = _parse_gemma4_args(args_str)
739 final_args_json = json.dumps(final_args, ensure_ascii=False)
740
677 - prev_streamed = self.streamed_args_for_tool[self.current_tool_id]
678 - if len(final_args_json) > len(prev_streamed):
679 - diff = final_args_json[len(prev_streamed) :]
680 - self.streamed_args_for_tool[self.current_tool_id] = final_args_json
681 - self.prev_tool_call_arr[self.current_tool_id]["arguments"] = final_args
682 -
683 - return DeltaMessage(
684 - tool_calls=[
741 + # The name is sent exactly once per call. We track that via the
742 + # per-call entry in prev_tool_call_arr (set either by the middle
743 + # path or by the batched-complete branch below), which is robust
744 + # even when several calls are drained in one delta.
745 + name_already_sent = bool(self.prev_tool_call_arr[idx].get("name"))
746 +
747 + if not name_already_sent:
748 + # Batched-complete call: emit id + name + full arguments once.
749 + self.streamed_args_for_tool[idx] = final_args_json
750 + self.prev_tool_call_arr[idx] = {
751 + "name": func_name,
752 + "arguments": final_args,
753 + }
754 + deltas.append(
755 + DeltaToolCall(
756 + index=idx,
757 + type="function",
758 + id=make_tool_call_id(),
759 + function=DeltaFunctionCall(
760 + name=func_name, arguments=final_args_json
761 + ).model_dump(exclude_none=True),
762 + )
763 + )
764 + else:
765 + # Incrementally-streamed call: flush the remaining argument
766 + # tail that was withheld during the middle phase.
767 + prev_streamed = self.streamed_args_for_tool[idx]
768 + if len(final_args_json) > len(prev_streamed):
769 + diff = final_args_json[len(prev_streamed) :]
770 + self.streamed_args_for_tool[idx] = final_args_json
771 + self.prev_tool_call_arr[idx]["arguments"] = final_args
772 + deltas.append(
773 DeltaToolCall(
686 - index=self.current_tool_id,
774 + index=idx,
775 function=DeltaFunctionCall(arguments=diff).model_dump(
776 exclude_none=True
777 ),
778 )
691 - ]
692 - )
779 + )
780 +
781 + # Advance streaming state past the calls completed in this delta. If a
782 + # further tool call is still being accumulated (start without a
783 + # matching end), point current_tool_id at it so the middle path can
784 + # stream its arguments next; otherwise settle on the last completed
785 + # call.
786 + if start_count > end_count:
787 + self.current_tool_id = end_count
788 + while len(self.prev_tool_call_arr) <= self.current_tool_id:
789 + self.prev_tool_call_arr.append({})
790 + self.streamed_args_for_tool.append("")
791 + self.current_tool_name_sent = bool(
792 + self.prev_tool_call_arr[self.current_tool_id].get("name")
793 + )
794 + else:
795 + self.current_tool_id = end_count - 1
796 + self.current_tool_name_sent = True
797
798 + if deltas:
799 + return DeltaMessage(tool_calls=deltas)
800 return None
801
802 def _emit_argument_diff(self, raw_args_str: str) -> DeltaMessage | None:
vllm/transformers_utils/config.py
+1
@@ -87,6 +87,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
87 ops_colqwen3="OpsColQwen3Config",
88 qwen3_vl_nemotron_embed="Qwen3VLNemotronEmbedConfig",
89 cosmos3_omni="Cosmos3Config",
90 + diffusion_gemma="DiffusionGemmaConfig",
91 deepseek_vl_v2="DeepseekVLV2Config",
92 deepseek_v32="DeepseekV3Config",
93 deepseek_v4="DeepseekV4Config",
vllm/transformers_utils/configs/__init__.py
+4
@@ -26,6 +26,8 @@ _CLASS_TO_MODULE: dict[str, str] = {
26 "OpsColQwen3Config": "vllm.transformers_utils.configs.colqwen3",
27 "Qwen3VLNemotronEmbedConfig": "vllm.transformers_utils.configs.colqwen3",
28 "Cosmos3Config": "vllm.transformers_utils.configs.cosmos3",
29 + "DiffusionGemmaConfig": "vllm.transformers_utils.configs.diffusion_gemma",
30 + "DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma",
31 "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2",
32 "DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4",
33 "DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr",
@@ -97,6 +99,8 @@ __all__ = [
99 "OpsColQwen3Config",
100 "Qwen3VLNemotronEmbedConfig",
101 "Cosmos3Config",
102 + "DiffusionGemmaConfig",
103 + "DiffusionGemmaTextConfig",
104 "DeepseekVLV2Config",
105 "DeepseekV3Config",
106 "DeepseekV4Config",
vllm/transformers_utils/configs/diffusion_gemma.py new
+44
@@ -0,0 +1,44 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 +from transformers import PretrainedConfig
6 +from transformers.models.gemma4.configuration_gemma4 import Gemma4VisionConfig
7 +
8 +
9 +def _init_text_config(self: PretrainedConfig, **kwargs: Any) -> None:
10 + PretrainedConfig.__init__(self, **kwargs)
11 + # DiffusionGemma always uses MoE and K=V sharing for full_attention
12 + # layers. The HF reference removed these config fields entirely.
13 + if getattr(self, "num_experts", None):
14 + self.enable_moe_block = True
15 + self.attention_k_eq_v = True
16 +
17 +
18 +class DiffusionGemmaTextConfig(PretrainedConfig):
19 + model_type = "diffusion_gemma_text"
20 +
21 + def __init__(self, **kwargs: Any):
22 + _init_text_config(self, **kwargs)
23 +
24 +
25 +class DiffusionGemmaConfig(PretrainedConfig):
26 + model_type = "diffusion_gemma"
27 +
28 + def __init__(
29 + self,
30 + text_config: dict[str, Any] | None = None,
31 + canvas_length: int = 256,
32 + self_conditioning_size: int | None = None,
33 + **kwargs: Any,
34 + ):
35 + self.text_config = DiffusionGemmaTextConfig(**(text_config or {}))
36 + self.canvas_length = canvas_length
37 + self.self_conditioning_size = self_conditioning_size
38 + vision_config = kwargs.pop("vision_config", None)
39 + if isinstance(vision_config, dict):
40 + self.vision_config = Gemma4VisionConfig(**vision_config)
41 + else:
42 + self.vision_config = vision_config
43 + self.audio_config = None
44 + PretrainedConfig.__init__(self, **kwargs)
vllm/transformers_utils/model_arch_config_convertor.py
+1
@@ -582,6 +582,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
582 "cohere_asr": CohereAsrModelArchConfigConvertor,
583 "dbrx": DbrxModelArchConfigConvertor,
584 "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor,
585 + "diffusion_gemma_text": Gemma4ModelArchConfigConvertor,
586 "ernie_mtp": ErnieMTPModelArchConfigConvertor,
587 "falcon": FalconModelArchConfigConvertor,
588 "falcon_mamba": MambaModelArchConfigConvertor,
vllm/v1/attention/backend.py
+4 -2
@@ -387,7 +387,7 @@ class CommonAttentionMetadata:
387 block_table_tensor: torch.Tensor
388 slot_mapping: torch.Tensor
389
390 - causal: bool = True
390 + causal: bool | torch.Tensor = True
391
392 # Needed by FastPrefillAttentionBuilder
393 logits_indices_padded: torch.Tensor | None = None
@@ -497,7 +497,9 @@ class CommonAttentionMetadata:
497 max_seq_len=self.max_seq_len,
498 block_table_tensor=self.block_table_tensor[:num_actual_reqs],
499 slot_mapping=self.slot_mapping[:num_actual_tokens],
500 - causal=self.causal,
500 + causal=self.causal[:num_actual_reqs]
501 + if isinstance(self.causal, torch.Tensor)
502 + else self.causal,
503 logits_indices_padded=self.logits_indices_padded,
504 num_logits_indices=self.num_logits_indices,
505 encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens),
vllm/v1/attention/backends/fa_utils.py
+6
@@ -131,6 +131,12 @@ def get_flash_attn_version(
131 and head_size != head_size_v
132 ):
133 upgrade_reason = "Diff-KV with sinks"
134 + elif (
135 + vllm_config is not None
136 + and vllm_config.model_config is not None
137 + and vllm_config.model_config.is_diffusion
138 + ):
139 + upgrade_reason = "Per-sequence causal (dynamic_causal) requires FA4"
140 if upgrade_reason:
141 logger.info_once(
142 "%s: upgrading FlashAttention 3 -> 4",
vllm/v1/attention/backends/flash_attn.py
+35 -3
@@ -267,7 +267,7 @@ class FlashAttentionMetadata:
267 prefix_scheduler_metadata: torch.Tensor | None = None
268 max_num_splits: int = 0
269
270 - causal: bool = True
270 + causal: bool | torch.Tensor = True
271
272 # PrefixLM bidirectional ranges for multimodal tokens.
273 # Shape: (num_seqs, max_ranges, 2) int32, [start, end] per range.
@@ -570,6 +570,9 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
570 self.scheduler_metadata[n:] = 0
571 scheduler_metadata = self.scheduler_metadata[:n]
572
573 + if isinstance(causal, torch.Tensor) and causal.dtype != torch.int32:
574 + causal = causal.to(torch.int32)
575 +
576 attn_metadata = FlashAttentionMetadata(
577 num_actual_tokens=num_actual_tokens,
578 max_query_len=max_query_len,
@@ -824,18 +827,46 @@ class FlashAttentionImpl(AttentionImpl):
827 if self.sliding_window is not None
828 else None
829 )
830 +
831 + causal = attn_metadata.causal
832 + is_dynamic_causal = isinstance(causal, torch.Tensor)
833 +
834 + # For non-causal (bidirectional) attention, make the
835 + # sliding window symmetric so queries attend in both
836 + # directions.
837 + if (
838 + sliding_window_size is not None
839 + and sliding_window_size[1] == 0
840 + and (is_dynamic_causal or causal is False)
841 + ):
842 + sliding_window_size = [
843 + sliding_window_size[0],
844 + sliding_window_size[0],
845 + ]
846 +
847 mm_prefix_ranges = attn_metadata.mm_prefix_range_tensor
848 mm_mask_mod = None
849 mm_aux = None
850 if (
851 mm_prefix_ranges is not None
832 - and attn_metadata.causal
852 + and not is_dynamic_causal
853 + and causal is True
854 and self.vllm_flash_attn_version == 4
855 ):
856 max_ranges = mm_prefix_ranges.shape[1]
857 mm_mask_mod = _make_mm_prefix_mask_mod(max_ranges)
858 mm_aux = [mm_prefix_ranges]
859
860 + dynamic_causal = None
861 + if isinstance(causal, torch.Tensor):
862 + if self.vllm_flash_attn_version != 4:
863 + raise NotImplementedError(
864 + "Per-sequence causal requires FA4. Current version: "
865 + f"FA{self.vllm_flash_attn_version}"
866 + )
867 + dynamic_causal = causal
868 + causal = False
869 +
870 flash_attn_varlen_func(
871 q=query[:num_actual_tokens],
872 k=key_cache,
@@ -846,7 +877,7 @@ class FlashAttentionImpl(AttentionImpl):
877 seqused_k=seqused_k,
878 max_seqlen_k=max_seqlen_k,
879 softmax_scale=self.scale,
849 - causal=attn_metadata.causal,
880 + causal=causal,
881 alibi_slopes=self.alibi_slopes,
882 window_size=sliding_window_size,
883 block_table=block_table,
@@ -856,6 +887,7 @@ class FlashAttentionImpl(AttentionImpl):
887 q_descale=q_descale,
888 k_descale=k_descale,
889 v_descale=v_descale,
890 + dynamic_causal=dynamic_causal,
891 num_splits=attn_metadata.max_num_splits,
892 s_aux=self.sinks,
893 mask_mod=mm_mask_mod,
vllm/v1/attention/backends/triton_attn.py
+8 -1
@@ -79,6 +79,8 @@ class TritonAttentionMetadata:
79 softmax_segm_max: torch.Tensor
80 softmax_segm_expsum: torch.Tensor
81
82 + causal: bool | torch.Tensor
83 +
84 # For cascade attention.
85 use_cascade: bool
86 common_prefix_len: int
@@ -219,6 +221,7 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
221 seq_lens=seq_lens,
222 block_table=block_table_tensor,
223 slot_mapping=slot_mapping,
224 + causal=common_attn_metadata.causal,
225 use_cascade=use_cascade,
226 common_prefix_len=common_prefix_len,
227 cu_prefix_query_lens=cu_prefix_query_lens,
@@ -271,6 +274,10 @@ class TritonAttentionBackend(AttentionBackend):
274
275 forward_includes_kv_cache_update: bool = False
276
277 + @classmethod
278 + def supports_non_causal(cls) -> bool:
279 + return True
280 +
281 @staticmethod
282 def get_name() -> str:
283 return "TRITON_ATTN"
@@ -619,7 +626,7 @@ class TritonAttentionImpl(AttentionImpl):
626 seqused_k=seqused_k,
627 max_seqlen_k=max_seqlen_k,
628 softmax_scale=self.scale,
622 - causal=True,
629 + causal=attn_metadata.causal,
630 alibi_slopes=self.alibi_slopes,
631 use_alibi_sqrt=self.use_alibi_sqrt,
632 window_size=self.sliding_window,
vllm/v1/attention/ops/triton_attention_helpers.py
+46 -11
@@ -153,6 +153,8 @@ def compute_tile_loop_bounds(
153 SLIDING_WINDOW: tl.constexpr,
154 USE_MM_PREFIX: tl.constexpr,
155 IS_3D: tl.constexpr,
156 + USE_CAUSAL: tl.constexpr = True,
157 + USE_PER_SEQ_CAUSAL: tl.constexpr = False,
158 CHUNK_LOOKBACK: tl.constexpr = -1,
159 CHUNK_SIZE: tl.constexpr = -1,
160 ):
@@ -163,10 +165,11 @@ def compute_tile_loop_bounds(
165
166 1. Longest prefix spanned by any query token in this q-block.
167 Clamped to ``seq_len`` (causal) or extended to it when
166 - mm_prefix is active (bidirectional ranges can reach past the
167 - causal prefix).
168 + mm_prefix is active or non-causal sequences need the full
169 + sequence.
170 2. Sliding-window pruning: narrows ``[tile_start, tile_end)`` to
171 only tiles that can contain an allowed key under SWA.
172 + For non-causal sequences, the window extends in both directions.
173 3. 3D scoping: when ``IS_3D`` is True, further narrows to the
174 segment's slice via ``(segm_idx * tiles_per_segment,
175 (segm_idx + 1) * tiles_per_segment)``.
@@ -179,9 +182,10 @@ def compute_tile_loop_bounds(
182 + (BLOCK_M - 1) // num_queries_per_kv
183 + 1
184 )
182 - if USE_MM_PREFIX:
183 - # image bidirectional attention ranges require a full range
184 - # including q_block padding to make sure doc mask is correct
185 + if USE_MM_PREFIX or USE_PER_SEQ_CAUSAL or (not USE_CAUSAL):
186 + # Non-causal or mixed batches need the full sequence range.
187 + # Per-element masking in compute_kv_seq_mask handles the
188 + # actual causal/non-causal boundary per sequence.
189 max_seq_prefix_len = tl.maximum(max_seq_prefix_len, seq_len)
190 else:
191 max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)
@@ -207,12 +211,17 @@ def compute_tile_loop_bounds(
211 # [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi]
212 q_abs = context_len + qpos_lo
213 if CHUNK_LOOKBACK > -1:
210 - # Chunked attention: align lower bound to the start of the
211 - # lookback'th previous chunk.
214 first_allowed_key = ((q_abs // CHUNK_SIZE) - CHUNK_LOOKBACK) * CHUNK_SIZE
215 else:
216 first_allowed_key = q_abs - SLIDING_WINDOW + 1
215 - last_allowed_key = context_len + qpos_hi
217 + if USE_PER_SEQ_CAUSAL or (not USE_CAUSAL):
218 + # Non-causal: keys can be AHEAD of query within the window
219 + last_allowed_key = tl.minimum(
220 + context_len + qpos_hi + SLIDING_WINDOW - 1,
221 + seq_len - 1,
222 + )
223 + else:
224 + last_allowed_key = context_len + qpos_hi
225 # Convert to tile indices and clamp
226 tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE)
227 tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles)
@@ -262,10 +271,14 @@ def compute_kv_seq_mask(
271 query_abs_pos,
272 seq_offset,
273 seq_idx,
274 + seq_len,
275 mm_prefix_range_ptr,
276 SLIDING_WINDOW: tl.constexpr,
277 USE_MM_PREFIX: tl.constexpr,
278 MAX_MM_RANGES: tl.constexpr,
279 + USE_CAUSAL: tl.constexpr = True,
280 + USE_PER_SEQ_CAUSAL: tl.constexpr = False,
281 + per_seq_causal_ptr=None,
282 CHUNK_LOOKBACK: tl.constexpr = -1,
283 CHUNK_SIZE: tl.constexpr = -1,
284 ):
@@ -279,9 +292,23 @@ def compute_kv_seq_mask(
292 Chunked attention takes precedence over sliding window when both
293 are non-default — the launcher zeros ``CHUNK_LOOKBACK`` whenever
294 sliding window is disabled.
295 +
296 + When ``USE_PER_SEQ_CAUSAL`` is set, each sequence carries its own
297 + causal flag via ``per_seq_causal_ptr``; non-causal sequences use a
298 + simple ``key < seq_len`` bound instead. ``USE_CAUSAL=False``
299 + disables causal masking entirely.
300 """
283 - # Compute attention mask: causal by default (key <= query)
284 - seq_mask = seq_offset[None, :] <= query_abs_pos
301 + if USE_PER_SEQ_CAUSAL:
302 + is_causal = tl.load(per_seq_causal_ptr + seq_idx)
303 + seq_mask = tl.where(
304 + is_causal,
305 + seq_offset[None, :] <= query_abs_pos,
306 + seq_offset[None, :] < seq_len,
307 + )
308 + elif USE_CAUSAL:
309 + seq_mask = seq_offset[None, :] <= query_abs_pos
310 + else:
311 + seq_mask = seq_offset[None, :] < seq_len
312
313 # Apply sliding window / chunked attention to base mask
314 # BEFORE mm_prefix OR.
@@ -293,7 +320,15 @@ def compute_kv_seq_mask(
320 <= CHUNK_LOOKBACK
321 )
322 elif SLIDING_WINDOW > 0:
296 - seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW)
323 + sw_left = (query_abs_pos - seq_offset) < SLIDING_WINDOW
324 + if USE_PER_SEQ_CAUSAL:
325 + sw_right = (seq_offset[None, :] - query_abs_pos) < SLIDING_WINDOW
326 + seq_mask = seq_mask & tl.where(is_causal, sw_left, sw_left & sw_right)
327 + elif not USE_CAUSAL:
328 + sw_right = (seq_offset[None, :] - query_abs_pos) < SLIDING_WINDOW
329 + seq_mask = seq_mask & sw_left & sw_right
330 + else:
331 + seq_mask = seq_mask & sw_left
332
333 # PrefixLM: extend mask with bidirectional ranges for multimodal tokens.
334 # Applied AFTER sliding window so mm_prefix ranges override SW restriction.
vllm/v1/attention/ops/triton_unified_attention.py
+63 -7
@@ -215,6 +215,9 @@ def kernel_unified_attention(
215 USE_SOFTCAP: tl.constexpr, # bool
216 USE_SINKS: tl.constexpr, # bool
217 SLIDING_WINDOW: tl.constexpr, # int
218 + USE_CAUSAL: tl.constexpr, # bool
219 + USE_PER_SEQ_CAUSAL: tl.constexpr, # bool
220 + per_seq_causal_ptr, # [num_seqs] bool, or None
221 USE_MM_PREFIX: tl.constexpr, # bool
222 MAX_MM_RANGES: tl.constexpr, # int
223 mm_prefix_range_ptr,
@@ -389,6 +392,8 @@ def kernel_unified_attention(
392 SLIDING_WINDOW,
393 USE_MM_PREFIX,
394 IS_3D,
395 + USE_CAUSAL,
396 + USE_PER_SEQ_CAUSAL,
397 CHUNK_LOOKBACK,
398 CHUNK_SIZE,
399 )
@@ -493,10 +498,14 @@ def kernel_unified_attention(
498 query_abs_pos,
499 seq_offset,
500 seq_idx,
501 + seq_len,
502 mm_prefix_range_ptr,
503 SLIDING_WINDOW,
504 USE_MM_PREFIX,
505 MAX_MM_RANGES,
506 + USE_CAUSAL,
507 + USE_PER_SEQ_CAUSAL,
508 + per_seq_causal_ptr,
509 CHUNK_LOOKBACK,
510 CHUNK_SIZE,
511 )
@@ -532,11 +541,19 @@ def kernel_unified_attention(
541
542 if SLIDING_WINDOW:
543 qpos_lo = q_block_local_idx * BLOCK_Q
535 - V = tl.where(
536 - (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW,
537 - V,
538 - 0.0,
539 - )
544 + dist = context_len + qpos_lo - seq_offset[:, None]
545 + if USE_PER_SEQ_CAUSAL:
546 + is_causal_seq = tl.load(per_seq_causal_ptr + seq_idx)
547 + sw_mask_v = tl.where(
548 + is_causal_seq,
549 + dist < SLIDING_WINDOW,
550 + (dist < SLIDING_WINDOW) & (dist > -SLIDING_WINDOW),
551 + )
552 + elif USE_CAUSAL:
553 + sw_mask_v = dist < SLIDING_WINDOW
554 + else:
555 + sw_mask_v = (dist < SLIDING_WINDOW) & (dist > -SLIDING_WINDOW)
556 + V = tl.where(sw_mask_v, V, 0.0)
557 if USE_PER_TOKEN_HEAD_SCALES:
558 # Per-token-head quant: apply v_scale to P instead of V.
559 P_v = (P * v_token_head_scales[None, :]).to(V.dtype)
@@ -802,7 +819,11 @@ def unified_attention(
819 # disabling this flag costs nothing.
820 use_td: bool = False,
821 ):
805 - assert causal, "Only causal attention is supported"
822 + # Resolve causal: bool or per-seq tensor.
823 + use_per_seq_causal = isinstance(causal, torch.Tensor)
824 + use_causal = bool(causal) if not use_per_seq_causal else True
825 + per_seq_causal_ptr = causal if use_per_seq_causal else None
826 +
827 if sinks is not None:
828 assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"
829
@@ -841,6 +862,26 @@ def unified_attention(
862 )
863 BLOCK_Q = BLOCK_M // num_queries_per_kv
864
865 + # Tuned launch parameters; ``None`` lets Triton pick its defaults.
866 + launch_num_warps: int | None = None
867 + launch_num_stages: int | None = None
868 +
869 + # head_size 256 with many query rows per sequence (e.g. diffusion-gemma
870 + # bidirectional canvas passes) is prefill-shaped, but the decode-oriented
871 + # defaults (BLOCK_Q=8, TILE=32, 4 warps) under-tile it. A wider KV tile +
872 + # more query rows per block + 8 warps is ~2x faster on B200.
873 + tuned_large_head = (
874 + head_size == 256
875 + and max_seqlen_q > 1
876 + and num_queries_per_kv <= 16
877 + and current_platform.is_device_capability_family(100)
878 + )
879 + if tuned_large_head:
880 + BLOCK_M = 32
881 + BLOCK_Q = BLOCK_M // num_queries_per_kv
882 + launch_num_warps = 8
883 + launch_num_stages = 2
884 +
885 # Ideally we would launch with kernel with:
886 # \sum_i[ceil(query_len[i] / BLOCK_Q)] blocks.
887 # However, it is slow to realize the query_lens on cpu.
@@ -869,6 +910,11 @@ def unified_attention(
910 head_size, sliding_window_val, q.element_size(), is_prefill=False
911 )
912
913 + # Wider KV tile for the tuned large-head path (see above). Only the 2D
914 + # path (used when max_seqlen_q > 1) reads TILE_SIZE_PREFILL.
915 + if tuned_large_head:
916 + TILE_SIZE_PREFILL = 128
917 +
918 # USE_TD requires BLOCK_SIZE % TILE_SIZE == 0 (enforced by a
919 # ``tl.static_assert`` in the kernel). The default prefill tile
920 # size (32) is larger than a common ``block_size=16``, so clamp it
@@ -964,6 +1010,12 @@ def unified_attention(
1010 grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments)
1011 tile_size = TILE_SIZE_DECODE
1012
1013 + launch_kwargs: dict[str, int] = {}
1014 + if launch_num_warps is not None:
1015 + launch_kwargs["num_warps"] = launch_num_warps
1016 + if launch_num_stages is not None:
1017 + launch_kwargs["num_stages"] = launch_num_stages
1018 +
1019 kernel_unified_attention[grid](
1020 output_ptr=out,
1021 segm_output_ptr=segm_output_ptr,
@@ -1002,10 +1054,13 @@ def unified_attention(
1054 USE_QQ_BIAS=use_qq_bias,
1055 USE_SOFTCAP=(softcap > 0),
1056 USE_SINKS=(sinks is not None),
1057 + SLIDING_WINDOW=(1 + window_size[0]),
1058 + USE_CAUSAL=use_causal,
1059 + USE_PER_SEQ_CAUSAL=use_per_seq_causal,
1060 + per_seq_causal_ptr=per_seq_causal_ptr,
1061 USE_MM_PREFIX=use_mm_prefix,
1062 MAX_MM_RANGES=max_mm_ranges,
1063 mm_prefix_range_ptr=mm_prefix_range,
1008 - SLIDING_WINDOW=(1 + window_size[0]),
1064 stride_k_cache_0=k.stride(0),
1065 stride_k_cache_1=k.stride(1),
1066 stride_k_cache_2=k.stride(2),
@@ -1033,6 +1088,7 @@ def unified_attention(
1088 CHUNK_SIZE=chunk_size,
1089 USE_TD=use_td,
1090 USE_TD_QO=use_td_qo,
1091 + **launch_kwargs,
1092 )
1093
1094 if use_3d:
vllm/v1/attention/ops/triton_unified_attention_diffkv.py
+1
@@ -226,6 +226,7 @@ def kernel_unified_attention_diffkv(
226 query_abs_pos,
227 seq_offset,
228 seq_idx,
229 + seq_len,
230 None, # mm_prefix_range_ptr
231 SLIDING_WINDOW,
232 False, # USE_MM_PREFIX
vllm/v1/core/sched/async_scheduler.py
+7 -3
@@ -27,10 +27,14 @@ class AsyncScheduler(Scheduler):
27 scheduler_output.pending_structured_output_tokens |= (
28 request.use_structured_output and request.num_output_placeholders > 0
29 )
30 - # The request will generate a new token plus num_spec_tokens
31 - # in this scheduling step.
30 + # The request will generate num_sampled_tokens_per_step new tokens
31 + # plus num_spec_tokens in this scheduling step. Diffusion has no AR
32 + # bonus token (num_sampled_tokens_per_step == 0) — only the canvas
33 + # (spec) tokens.
34 cur_num_spec_tokens = len(spec_decode_tokens.get(req_id, ()))
33 - request.num_output_placeholders += 1 + cur_num_spec_tokens
35 + request.num_output_placeholders += (
36 + self.num_sampled_tokens_per_step + cur_num_spec_tokens
37 + )
38 # Add placeholders for the new draft/spec tokens.
39 # We will update the actual spec token ids in the worker process.
40 request.spec_token_ids = self._spec_token_placeholders
vllm/v1/core/sched/scheduler.py
+16 -6
@@ -113,6 +113,10 @@ class Scheduler(SchedulerInterface):
113 self.kv_events_config is not None
114 and self.kv_events_config.enable_kv_cache_events
115 )
116 + # Diffusion models may not sample any tokens for a denoising step.
117 + self.num_sampled_tokens_per_step = (
118 + 1 if not vllm_config.model_config.is_diffusion else 0
119 + )
120
121 # Create KVConnector for the Scheduler. Note that each Worker
122 # will have a corresponding KVConnector with Role=WORKER.
@@ -212,9 +216,9 @@ class Scheduler(SchedulerInterface):
216
217 speculative_config = vllm_config.speculative_config
218 self.use_eagle = False
215 - self.num_spec_tokens = self.num_lookahead_tokens = 0
216 - if speculative_config:
217 - self.num_spec_tokens = speculative_config.num_speculative_tokens
219 + self.num_spec_tokens = vllm_config.num_speculative_tokens
220 + self.num_lookahead_tokens = 0
221 + if speculative_config is not None:
222 if speculative_config.use_eagle():
223 self.use_eagle = True
224 self.num_lookahead_tokens = self.num_spec_tokens
@@ -425,7 +429,10 @@ class Scheduler(SchedulerInterface):
429 # Make sure the input position does not exceed the max model len.
430 # This is necessary when using spec decoding.
431 num_new_tokens = min(
428 - num_new_tokens, self.max_model_len - 1 - request.num_computed_tokens
432 + num_new_tokens,
433 + self.max_model_len
434 + - request.num_computed_tokens
435 + - self.num_sampled_tokens_per_step,
436 )
437
438 # Schedule encoder inputs.
@@ -1473,9 +1480,12 @@ class Scheduler(SchedulerInterface):
1480 scheduled_spec_token_ids = (
1481 scheduler_output.scheduled_spec_decode_tokens.get(req_id)
1482 )
1476 - if scheduled_spec_token_ids and generated_token_ids:
1483 + if scheduled_spec_token_ids and (
1484 + generated_token_ids or self.num_sampled_tokens_per_step == 0
1485 + ):
1486 num_draft_tokens = len(scheduled_spec_token_ids)
1478 - num_accepted = len(generated_token_ids) - 1
1487 + num_sampled = self.num_sampled_tokens_per_step
1488 + num_accepted = max(len(generated_token_ids) - num_sampled, 0)
1489 num_rejected = num_draft_tokens - num_accepted
1490 # num_computed_tokens represents the number of tokens
1491 # processed in the current step, considering scheduled
vllm/v1/cudagraph_dispatcher.py
+1 -5
@@ -34,11 +34,7 @@ class CudagraphDispatcher:
34 def __init__(self, vllm_config: VllmConfig):
35 self.vllm_config = vllm_config
36 self.compilation_config = vllm_config.compilation_config
37 - self.uniform_decode_query_len = (
38 - 1
39 - if not self.vllm_config.speculative_config
40 - else 1 + self.vllm_config.speculative_config.num_speculative_tokens
41 - )
37 + self.uniform_decode_query_len = 1 + self.vllm_config.num_speculative_tokens
38
39 # Dict to store valid cudagraph dispatching keys.
40 self.cudagraph_keys: dict[CUDAGraphMode, set[BatchDescriptor]] = {
vllm/v1/engine/core.py
+14 -13
@@ -156,6 +156,9 @@ class EngineCore:
156 hash_block_size=hash_block_size,
157 )
158 self.use_spec_decode = vllm_config.speculative_config is not None
159 + self.check_for_draft_tokens = (
160 + self.use_spec_decode or vllm_config.model_config.is_diffusion
161 + )
162 if self.scheduler.connector is not None: # type: ignore
163 self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore
164
@@ -475,8 +478,7 @@ class EngineCore:
478 # When using async scheduling we can't get draft token ids in advance,
479 # so we update draft token ids in the worker process and don't
480 # need to update draft token ids here.
478 - if not self.async_scheduling and self.use_spec_decode and model_executed:
479 - # Take the draft token ids.
481 + if self.check_for_draft_tokens and not self.async_scheduling and model_executed:
482 draft_token_ids = self.model_executor.take_draft_token_ids()
483 if draft_token_ids is not None:
484 self.scheduler.update_draft_token_ids(draft_token_ids)
@@ -575,18 +577,17 @@ class EngineCore:
577 # in a field and do it immediately once step_with_batch_queue is
578 # re-called. The latter slightly favors TTFT over TPOT/throughput.
579 if deferred_scheduler_output:
578 - # If we are doing speculative decoding with structured output,
579 - # we need to get the draft token ids from the prior step before
580 - # we can compute the grammar bitmask for the deferred request.
581 - if self.use_spec_decode:
580 + # When draft tokens are used with structured output, validate them
581 + # before computing the grammar bitmask for the deferred request.
582 + if self.check_for_draft_tokens:
583 draft_token_ids = self.model_executor.take_draft_token_ids()
583 - assert draft_token_ids is not None
584 - # Update the draft token ids in the scheduler output to
585 - # filter out the invalid spec tokens, which will be padded
586 - # with -1 and skipped by the grammar bitmask computation.
587 - self.scheduler.update_draft_token_ids_in_output(
588 - draft_token_ids, deferred_scheduler_output
589 - )
584 + if draft_token_ids is not None:
585 + # Update the draft token ids in the scheduler output to
586 + # filter out the invalid spec tokens, which will be padded
587 + # with -1 and skipped by the grammar bitmask computation.
588 + self.scheduler.update_draft_token_ids_in_output(
589 + draft_token_ids, deferred_scheduler_output
590 + )
591 # We now have the tokens needed to compute the bitmask for the
592 # deferred request. Get the bitmask and call sample tokens.
593 grammar_output = self.scheduler.get_grammar_bitmask(
vllm/v1/metrics/loggers.py
+7 -2
@@ -110,7 +110,9 @@ class LoggingStatLogger(StatLoggerBase):
110 self.connector_prefix_caching_metrics = CachingMetrics()
111 self.mm_caching_metrics = CachingMetrics()
112
113 - self.spec_decoding_logging = SpecDecodingLogging()
113 + model_config = self.vllm_config.model_config
114 + is_diffusion = model_config is not None and model_config.is_diffusion
115 + self.spec_decoding_logging = SpecDecodingLogging(is_diffusion=is_diffusion)
116 kv_transfer_config = self.vllm_config.kv_transfer_config
117 self.kv_connector_logging = KVConnectorLogging(kv_transfer_config)
118 self.cudagraph_logging = None
@@ -436,7 +438,10 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
438 per_engine_labelvalues = self.per_engine_labelvalues
439
440 self.spec_decoding_prom = self._spec_decoding_cls(
439 - vllm_config.speculative_config, labelnames, per_engine_labelvalues
441 + vllm_config.speculative_config,
442 + labelnames,
443 + per_engine_labelvalues,
444 + is_diffusion=vllm_config.model_config.is_diffusion,
445 )
446 self.kv_connector_prom = self._kv_connector_cls(
447 vllm_config, labelnames, per_engine_labelvalues
vllm/v1/spec_decode/metrics.py
+107 -45
@@ -53,7 +53,11 @@ class SpecDecodingLogging:
53 before resetting to zero.
54 """
55
56 - def __init__(self):
56 + def __init__(self, is_diffusion: bool = False):
57 + # Diffusion (dLLM) models reuse the spec-decode data path with
58 + # overloaded semantics, so the raw spec-decode framing (drafts, bonus
59 + # token, per-position vector) is logged with diffusion-native terms.
60 + self.is_diffusion = is_diffusion
61 self.reset()
62
63 def reset(self):
@@ -85,6 +89,17 @@ class SpecDecodingLogging:
89 draft_throughput = num_draft_tokens / elapsed_time
90 accepted_throughput = num_accepted_tokens / elapsed_time
91
92 + if self.is_diffusion:
93 + self._log_diffusion(
94 + log_fn,
95 + num_denoising_steps=num_drafts,
96 + num_canvas_tokens=num_draft_tokens,
97 + num_committed_tokens=num_accepted_tokens,
98 + committed_throughput=accepted_throughput,
99 + )
100 + self.reset()
101 + return
102 +
103 draft_acceptance_rate = (
104 num_accepted_tokens / num_draft_tokens * 100
105 if num_draft_tokens > 0
@@ -117,6 +132,43 @@ class SpecDecodingLogging:
132 )
133 self.reset()
134
135 + def _log_diffusion(
136 + self,
137 + log_fn,
138 + num_denoising_steps: int,
139 + num_canvas_tokens: int,
140 + num_committed_tokens: int,
141 + committed_throughput: float,
142 + ):
143 + # Each "draft" is one denoising step that re-evaluates the canvas block
144 + # and finalizes some of its positions.
145 + mean_committed_per_step = (
146 + num_committed_tokens / num_denoising_steps
147 + if num_denoising_steps > 0
148 + else float("nan")
149 + )
150 + mean_steps_per_canvas = (
151 + num_canvas_tokens / num_committed_tokens
152 + if num_committed_tokens > 0
153 + else float("nan")
154 + )
155 +
156 + log_fn(
157 + "DiffusionDecoding metrics: "
158 + "Committed token throughput: %.2f tokens/s, "
159 + "Mean denoising steps per canvas: %.2f, "
160 + "Mean tokens committed per denoising step: %.2f, "
161 + "Committed: %d tokens, "
162 + "Denoising steps: %d, "
163 + "Canvas positions evaluated: %d",
164 + committed_throughput,
165 + mean_steps_per_canvas,
166 + mean_committed_per_step,
167 + num_committed_tokens,
168 + num_denoising_steps,
169 + num_canvas_tokens,
170 + )
171 +
172
173 class SpecDecodingProm:
174 """Record spec decoding metrics in Prometheus.
@@ -146,56 +198,66 @@ class SpecDecodingProm:
198 speculative_config: SpeculativeConfig | None,
199 labelnames: list[str],
200 per_engine_labelvalues: dict[int, list[object]],
201 + is_diffusion: bool = False,
202 ):
150 - self.spec_decoding_enabled = speculative_config is not None
203 + # Diffusion (dLLM) models reuse the spec-decode counters but expose them
204 + # under diffusion-native names; the per-position acceptance vector does
205 + # not apply, so it is omitted.
206 + self.is_diffusion = is_diffusion
207 + self.spec_decoding_enabled = speculative_config is not None or is_diffusion
208 if not self.spec_decoding_enabled:
209 return
210
154 - counter_drafts = self._counter_cls(
155 - name="vllm:spec_decode_num_drafts",
156 - documentation="Number of spec decoding drafts.",
157 - labelnames=labelnames,
158 - )
159 - self.counter_spec_decode_num_drafts = create_metric_per_engine(
160 - counter_drafts, per_engine_labelvalues
161 - )
162 -
163 - counter_draft_tokens = self._counter_cls(
164 - name="vllm:spec_decode_num_draft_tokens",
165 - documentation="Number of draft tokens.",
166 - labelnames=labelnames,
167 - )
168 - self.counter_spec_decode_num_draft_tokens = create_metric_per_engine(
169 - counter_draft_tokens, per_engine_labelvalues
170 - )
211 + if is_diffusion:
212 + counter_specs = [
213 + ("vllm:diffusion_num_denoising_steps", "Number of denoising steps."),
214 + (
215 + "vllm:diffusion_num_canvas_positions",
216 + "Number of canvas positions evaluated.",
217 + ),
218 + (
219 + "vllm:diffusion_num_committed_tokens",
220 + "Number of committed (finalized) tokens.",
221 + ),
222 + ]
223 + else:
224 + counter_specs = [
225 + ("vllm:spec_decode_num_drafts", "Number of spec decoding drafts."),
226 + ("vllm:spec_decode_num_draft_tokens", "Number of draft tokens."),
227 + ("vllm:spec_decode_num_accepted_tokens", "Number of accepted tokens."),
228 + ]
229 +
230 + counters = [
231 + create_metric_per_engine(
232 + self._counter_cls(name=name, documentation=doc, labelnames=labelnames),
233 + per_engine_labelvalues,
234 + )
235 + for name, doc in counter_specs
236 + ]
237 + # num_drafts/num_draft_tokens/num_accepted_tokens map onto denoising
238 + # steps/canvas positions/committed tokens in the diffusion path.
239 + self.counter_spec_decode_num_drafts = counters[0]
240 + self.counter_spec_decode_num_draft_tokens = counters[1]
241 + self.counter_spec_decode_num_accepted_tokens = counters[2]
242
172 - counter_accepted_tokens = self._counter_cls(
173 - name="vllm:spec_decode_num_accepted_tokens",
174 - documentation="Number of accepted tokens.",
175 - labelnames=labelnames,
176 - )
177 - self.counter_spec_decode_num_accepted_tokens = create_metric_per_engine(
178 - counter_accepted_tokens, per_engine_labelvalues
179 - )
180 -
181 - assert speculative_config is not None
182 - num_spec_tokens = (
183 - speculative_config.num_speculative_tokens
184 - if self.spec_decoding_enabled
185 - else 0
186 - )
187 - pos_labelnames = labelnames + ["position"]
188 - base_counter = self._counter_cls(
189 - name="vllm:spec_decode_num_accepted_tokens_per_pos",
190 - documentation="Accepted tokens per draft position.",
191 - labelnames=pos_labelnames,
192 - )
243 self.counter_spec_decode_num_accepted_tokens_per_pos: dict[
244 int, list[prometheus_client.Counter]
195 - ] = {
196 - idx: [base_counter.labels(*lv, str(pos)) for pos in range(num_spec_tokens)]
197 - for idx, lv in per_engine_labelvalues.items()
198 - }
245 + ] = {}
246 + if not is_diffusion:
247 + assert speculative_config is not None
248 + num_spec_tokens = speculative_config.num_speculative_tokens
249 + pos_labelnames = labelnames + ["position"]
250 + base_counter = self._counter_cls(
251 + name="vllm:spec_decode_num_accepted_tokens_per_pos",
252 + documentation="Accepted tokens per draft position.",
253 + labelnames=pos_labelnames,
254 + )
255 + self.counter_spec_decode_num_accepted_tokens_per_pos = {
256 + idx: [
257 + base_counter.labels(*lv, str(pos)) for pos in range(num_spec_tokens)
258 + ]
259 + for idx, lv in per_engine_labelvalues.items()
260 + }
261
262 def observe(self, spec_decoding_stats: SpecDecodingStats, engine_idx: int = 0):
263 if not self.spec_decoding_enabled:
@@ -210,6 +272,6 @@ class SpecDecodingProm:
272 spec_decoding_stats.num_accepted_tokens
273 )
274 for pos, counter in enumerate(
213 - self.counter_spec_decode_num_accepted_tokens_per_pos[engine_idx]
275 + self.counter_spec_decode_num_accepted_tokens_per_pos.get(engine_idx, [])
276 ):
277 counter.inc(spec_decoding_stats.num_accepted_tokens_per_pos[pos])
vllm/v1/structured_output/__init__.py
+9 -6
@@ -211,11 +211,8 @@ class StructuredOutputManager:
211 if not structured_output_request_ids:
212 return None
213
214 - max_num_spec_tokens = 0
215 - if self.vllm_config.speculative_config is not None:
216 - max_num_spec_tokens = (
217 - self.vllm_config.speculative_config.num_speculative_tokens
218 - )
214 + # Covers both speculative decoding and diffusion LLMs (canvas_length).
215 + max_num_spec_tokens = self.vllm_config.num_speculative_tokens
216
217 if self._grammar_bitmask is None:
218 assert self.backend is not None
@@ -277,7 +274,13 @@ class StructuredOutputManager:
274
275 state_advancements = 0
276 req_tokens = scheduled_spec_decode_tokens.get(req_id, ())
280 - for token in itertools.chain(req_tokens, (-1,)):
277 + if self.vllm_config.model_config.is_diffusion and req_tokens:
278 + # Diffusion LLMs don't sample a bonus token after the
279 + # scheduled positions, so don't append the -1 placeholder.
280 + token_iter: Iterable[int] = req_tokens
281 + else:
282 + token_iter = itertools.chain(req_tokens, (-1,))
283 + for token in token_iter:
284 self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))
285 if token == -1:
286 # Stop advancing the grammar once we hit a padding token.
vllm/v1/worker/gpu/input_batch.py
+16 -7
@@ -302,6 +302,7 @@ def _combine_sampled_and_draft_tokens_kernel(
302 cu_num_logits_ptr,
303 logits_indices_ptr,
304 BLOCK_SIZE: tl.constexpr,
305 + NUM_NEW_SAMPLED_TOKENS: tl.constexpr = 1,
306 ):
307 batch_idx = tl.program_id(0)
308 req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
@@ -310,7 +311,7 @@ def _combine_sampled_and_draft_tokens_kernel(
311 cu_num_logits_start = tl.load(cu_num_logits_ptr + batch_idx)
312 cu_num_logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1)
313 num_logits = cu_num_logits_end - cu_num_logits_start
313 - num_draft_tokens = num_logits - 1
314 + num_draft_tokens = num_logits - NUM_NEW_SAMPLED_TOKENS
315
316 # Compute the logits indices.
317 block = tl.arange(0, BLOCK_SIZE)
@@ -328,9 +329,10 @@ def _combine_sampled_and_draft_tokens_kernel(
329 # Handling prefill tokens. No sampled or draft tokens.
330 return
331
331 - # Write the last sampled token ID to input_ids.
332 - last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx)
333 - tl.store(input_ids_ptr + query_end - num_logits, last_token_id)
332 + if NUM_NEW_SAMPLED_TOKENS > 0:
333 + # Write the last sampled token ID to input_ids.
334 + last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx)
335 + tl.store(input_ids_ptr + query_end - num_logits, last_token_id)
336
337 # Write the draft tokens (if any) to input_ids.
338 if num_draft_tokens > 0:
@@ -356,7 +358,11 @@ def combine_sampled_and_draft_tokens(
358 draft_tokens: torch.Tensor,
359 cu_num_logits: torch.Tensor,
360 num_logits: int,
361 + num_new_sampled_tokens: int = 1, # excl accepted draft tokens, a.k.a bonus tokens
362 ) -> torch.Tensor:
363 + assert num_new_sampled_tokens in (0, 1), (
364 + f"num_new_sampled_tokens must be 0 or 1, got {num_new_sampled_tokens}"
365 + )
366 # use idx_mapping.shape[0] for actual request count
367 num_reqs = idx_mapping.shape[0]
368 num_speculative_steps = draft_tokens.shape[-1]
@@ -377,9 +383,12 @@ def combine_sampled_and_draft_tokens(
383 draft_tokens.stride(0),
384 cu_num_logits,
385 logits_indices,
380 - # NOTE(woosuk): Add 1 to ensure the block can cover the last sampled token
381 - # in addition to all draft tokens.
382 - BLOCK_SIZE=triton.next_power_of_2(num_speculative_steps + 1),
386 + NUM_NEW_SAMPLED_TOKENS=num_new_sampled_tokens,
387 + # NOTE(woosuk): Add num_new_sampled_tokens to ensure the block covers the
388 + # last sampled token in addition to all draft tokens.
389 + BLOCK_SIZE=triton.next_power_of_2(
390 + num_speculative_steps + num_new_sampled_tokens
391 + ),
392 )
393 return logits_indices
394
vllm/v1/worker/gpu/model_runner.py
+56 -51
@@ -78,7 +78,6 @@ from vllm.v1.worker.gpu.input_batch import (
78 InputBuffers,
79 combine_sampled_and_draft_tokens,
80 expand_idx_mapping,
81 - get_num_sampled_and_rejected,
81 post_update,
82 post_update_num_computed_tokens,
83 prepare_pos_seq_lens,
@@ -185,11 +184,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
184
185 # Speculative decoding.
186 self.speculator = None
188 - self.num_speculative_steps = 0
187 self.use_aux_hidden_state_outputs = False
188 + self.num_speculative_steps = vllm_config.num_speculative_tokens
189 if self.speculative_config is not None:
191 - self.num_speculative_steps = self.speculative_config.num_speculative_tokens
192 -
190 if self.is_last_pp_rank:
191 self.speculator = init_speculator(self.vllm_config, self.device)
192
@@ -204,7 +201,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
201
202 # Draft tokens propagation - for spec-dec + struct outputs.
203 self.draft_tokens_handler = DraftTokensHandler(self.device)
207 - self.uniform_decode_query_len = 1 + self.num_speculative_steps
204
205 # Pooling models.
206 self.is_pooling_model = self.model_config.runner_type == "pooling"
@@ -232,38 +228,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
228 device=self.device,
229 )
230
231 + # Samplers and decode_query_len created in load_model() after
232 + # model_state exists (num_new_sampled_tokens_per_step from ModelState).
233 self.sampler: Sampler | None = None
234 self.rejection_sampler: RejectionSampler | None = None
235 self.prompt_logprobs_worker: PromptLogprobsWorker | None = None
236 self.structured_outputs_worker: StructuredOutputsWorker | None = None
239 - if self.is_last_pp_rank and not self.is_pooling_model:
240 - # Initialize sampling-related workers.
241 - # These components are only set up on the last PP rank and
242 - # for generative (non-pooling) models.
243 - self.sampler = Sampler(
244 - max_num_reqs=self.max_num_reqs,
245 - vocab_size=self.vocab_size,
246 - device=self.device,
247 - req_states=self.req_states,
248 - logprobs_mode=self.model_config.logprobs_mode,
249 - num_speculative_tokens=self.num_speculative_steps + 1,
250 - use_fp64_gumbel=self.model_config.use_fp64_gumbel,
251 - )
252 - if self.speculative_config is not None:
253 - self.rejection_sampler = RejectionSampler(
254 - self.sampler,
255 - self.speculative_config,
256 - self.device,
257 - )
258 - self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
259 - self.structured_outputs_worker = StructuredOutputsWorker(
260 - max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1),
261 - vocab_size=self.vocab_size,
262 - device=self.device,
263 - )
264 -
265 - # For CUDA graphs, and will init cudagraph_manager after init_attn_backend.
266 - self.decode_query_len = self.num_speculative_steps + 1
237 self.cudagraph_manager: ModelCudaGraphManager | None = None
238 # LoRA-related workers.
239 self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
@@ -335,6 +305,40 @@ class GPUModelRunner(LoRAModelRunnerMixin):
305 self.model_state = init_model_state(
306 self.vllm_config, self.model, self.encoder_cache, self.device
307 )
308 +
309 + self.decode_query_len = (
310 + self.num_speculative_steps
311 + + self.model_state.num_new_sampled_tokens_per_step
312 + )
313 +
314 + # Initialize samplers. Model states may override via custom_sampler().
315 + if self.is_last_pp_rank and not self.is_pooling_model:
316 + self.sampler = Sampler(
317 + max_num_reqs=self.max_num_reqs,
318 + vocab_size=self.vocab_size,
319 + device=self.device,
320 + req_states=self.req_states,
321 + logprobs_mode=self.model_config.logprobs_mode,
322 + num_speculative_tokens=self.decode_query_len,
323 + use_fp64_gumbel=self.model_config.use_fp64_gumbel,
324 + )
325 + custom = self.model_state.custom_sampler(self.sampler)
326 +
327 + if custom:
328 + self.sampler, self.rejection_sampler = custom
329 + elif self.speculative_config is not None:
330 + self.rejection_sampler = RejectionSampler(
331 + self.sampler,
332 + self.speculative_config,
333 + self.device,
334 + )
335 + self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
336 + self.structured_outputs_worker = StructuredOutputsWorker(
337 + max_num_logits=self.max_num_reqs * self.decode_query_len,
338 + vocab_size=self.vocab_size,
339 + device=self.device,
340 + )
341 +
342 if self.is_pooling_model and self.is_last_pp_rank:
343 self.pooling_runner = PoolingRunner(self.model)
344 eplb_models_added |= self.eplb.maybe_register_model(
@@ -447,7 +451,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
451 cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
452 attn_cg_support.min_cg_support,
453 attn_cg_support.min_cg_attn_backend,
450 - self.uniform_decode_query_len,
454 + self.decode_query_len,
455 self.parallel_config.tensor_parallel_size,
456 self.kv_cache_config,
457 self.max_num_reqs,
@@ -710,6 +714,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
714 return cuda_graph_size
715
716 def _remove_request(self, req_id: str) -> bool:
717 + # Call model_state.remove_request *before* req_states.remove_request
718 + # so the model_state can still look up the slot index.
719 + self.model_state.remove_request(req_id)
720 req_idx = self.req_states.remove_request(req_id)
721 if req_idx is None:
722 return False
@@ -857,16 +864,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
864 dtype=np.int32,
865 count=num_reqs,
866 )
867 + num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step
868 total_num_draft_tokens = int(num_draft_tokens_per_req.sum())
861 - total_num_logits = num_reqs + total_num_draft_tokens
862 -
863 - num_logits = num_draft_tokens_per_req + 1
869 + total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
870 + num_logits = num_draft_tokens_per_req + num_bonus_tokens
871 cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
872 cu_num_logits_np[0] = 0
873 np.cumsum(num_logits, out=cu_num_logits_np[1:])
874 cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
875
869 - max_expand_len = self.num_speculative_steps + 1
876 + max_expand_len = self.decode_query_len
877 expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
878 idx_mapping, total_num_logits, cu_num_logits, max_expand_len
879 )
@@ -935,6 +942,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
942 self.req_states.draft_tokens,
943 cu_num_logits,
944 total_num_logits,
945 + self.model_state.num_new_sampled_tokens_per_step,
946 )
947
948 # CPU upper bound on seq_lens; padded entries left at zero.
@@ -1027,8 +1035,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1035 grammar_output.grammar_bitmask,
1036 )
1037
1030 - if input_batch.num_draft_tokens == 0:
1031 - # No draft tokens (common case).
1038 + if input_batch.num_draft_tokens == 0 or self.rejection_sampler is None:
1039 assert self.sampler is not None
1040 sampler_output = self.sampler(logits, input_batch)
1041 else:
@@ -1042,16 +1049,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1049 self.speculator.draft_logits,
1050 )
1051
1045 - # Get the number of sampled and rejected tokens.
1046 - # For chunked prefills, num_sampled and num_rejected are both 0.
1047 - num_sampled, num_rejected = get_num_sampled_and_rejected(
1048 - sampler_output.num_sampled,
1049 - input_batch.seq_lens,
1050 - input_batch.cu_num_logits,
1051 - input_batch.idx_mapping,
1052 - self.req_states.prefill_len.gpu,
1053 - )
1054 - return sampler_output, num_sampled, num_rejected
1052 + return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected
1053
1054 def postprocess_sampled(
1055 self,
@@ -1448,7 +1446,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1446 mm_inputs=mm_inputs,
1447 )
1448 self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
1451 - self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)
1449 +
1450 + if self.num_speculative_steps > 0:
1451 + # Spec-decode and diffusion LLMs both use draft tokens but the latter does
1452 + # not have a speculator (i.e. self.speculator is None)
1453 + self.draft_tokens_handler.set_draft_tokens(
1454 + input_batch,
1455 + self.req_states.draft_tokens[input_batch.idx_mapping],
1456 + )
1457
1458 # Post-step KV connector related operations.
1459 kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
vllm/v1/worker/gpu/model_states/__init__.py
+5
@@ -13,6 +13,11 @@ def init_model_state(
13 encoder_cache: EncoderCache | None,
14 device: torch.device,
15 ):
16 + # Let the model provide its own ModelState if it defines one.
17 + if hasattr(model, "get_model_state_cls"):
18 + cls = model.get_model_state_cls()
19 + return cls(vllm_config, model, encoder_cache, device)
20 +
21 if (
22 "WhisperForConditionalGeneration" in vllm_config.model_config.architectures
23 or "CohereAsrForConditionalGeneration" in vllm_config.model_config.architectures
vllm/v1/worker/gpu/model_states/interface.py
+16
@@ -53,6 +53,9 @@ class ModelState(ABC):
53 def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
54 return None
55
56 + def remove_request(self, req_id: str) -> None:
57 + return None
58 +
59 def apply_staged_writes(self) -> None:
60 return None
61
@@ -89,3 +92,16 @@ class ModelState(ABC):
92 for_capture: bool = False,
93 ) -> dict[str, Any]:
94 raise NotImplementedError
95 +
96 + def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None:
97 + """Wrap or replace the default sampler.
98 +
99 + Called after model loading with the already-constructed base
100 + ``Sampler``. Return ``None`` to keep the defaults, or
101 + ``(sampler, rejection_sampler | None)`` to override.
102 + """
103 + return None
104 +
105 + num_new_sampled_tokens_per_step: int = 1
106 + """New tokens sampled on each decode step
107 + (excluding accepted draft tokens, a.k.a num bonus tokens)."""
vllm/v1/worker/gpu/sample/output.py
+1
@@ -13,3 +13,4 @@ class SamplerOutput:
13 logprobs_tensors: LogprobsTensors | None
14 num_nans: torch.Tensor | None
15 num_sampled: torch.Tensor | None
16 + num_rejected: torch.Tensor | None = None
vllm/v1/worker/gpu/sample/sampler.py
+15 -2
@@ -12,7 +12,7 @@ from vllm.v1.sample.ops.topk_topp_sampler import (
12 flashinfer_sample,
13 flashinfer_sampler_supported,
14 )
15 -from vllm.v1.worker.gpu.input_batch import InputBatch
15 +from vllm.v1.worker.gpu.input_batch import InputBatch, get_num_sampled_and_rejected
16 from vllm.v1.worker.gpu.metrics.logits import get_num_nans
17 from vllm.v1.worker.gpu.sample.bad_words import BadWordsState
18 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
@@ -44,6 +44,7 @@ class Sampler:
44 self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default.
45 self.use_fp64_gumbel = use_fp64_gumbel
46
47 + self.req_states = req_states
48 self.sampling_states = SamplingStates(max_num_reqs, vocab_size)
49 self.penalties_state = PenaltiesState(req_states)
50 self.logit_bias_state = LogitBiasState(max_num_reqs, device)
@@ -118,6 +119,17 @@ class Sampler:
119 else:
120 logprobs_tensors = None
121
122 + # 1 sampled token per request, except chunked-prefill requests
123 + # (seq_len < prefill_len) which aren't done prefilling and produce no
124 + # output token. num_rejected is always 0 here (one logit per request).
125 + num_sampled, num_rejected = get_num_sampled_and_rejected(
126 + input_batch.seq_lens.new_ones(input_batch.num_reqs),
127 + input_batch.seq_lens,
128 + input_batch.cu_num_logits,
129 + input_batch.idx_mapping,
130 + self.req_states.prefill_len.gpu,
131 + )
132 +
133 # These are GPU tensors.
134 sampler_output = SamplerOutput(
135 # The sampled tokens are expanded to 2D tensor with shape
@@ -126,7 +138,8 @@ class Sampler:
138 sampled_token_ids=sampled.view(-1, 1),
139 logprobs_tensors=logprobs_tensors,
140 num_nans=num_nans,
129 - num_sampled=input_batch.seq_lens.new_ones(input_batch.num_reqs),
141 + num_sampled=num_sampled,
142 + num_rejected=num_rejected,
143 )
144 return sampler_output
145
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+13 -1
@@ -6,7 +6,10 @@ from vllm.config import SpeculativeConfig
6 from vllm.triton_utils import tl, triton
7 from vllm.v1.outputs import LogprobsTensors
8 from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates
9 -from vllm.v1.worker.gpu.input_batch import InputBatch
9 +from vllm.v1.worker.gpu.input_batch import (
10 + InputBatch,
11 + get_num_sampled_and_rejected,
12 +)
13 from vllm.v1.worker.gpu.metrics.logits import get_num_nans
14 from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
15 from vllm.v1.worker.gpu.sample.output import SamplerOutput
@@ -136,9 +139,18 @@ class RejectionSampler:
139 else logits,
140 )
141
142 + num_sampled, num_rejected = get_num_sampled_and_rejected(
143 + num_sampled,
144 + input_batch.seq_lens,
145 + input_batch.cu_num_logits,
146 + input_batch.idx_mapping,
147 + self.sampler.req_states.prefill_len.gpu,
148 + )
149 +
150 return SamplerOutput(
151 sampled_token_ids=sampled,
152 logprobs_tensors=logprobs_tensors,
153 num_nans=num_nans,
154 num_sampled=num_sampled,
155 + num_rejected=num_rejected,
156 )
vllm/v1/worker/gpu/spec_decode/utils.py
+4
@@ -35,6 +35,10 @@ class DraftTokensHandler:
35 self.copy_stream.wait_stream(current_stream)
36 with torch.cuda.stream(self.copy_stream):
37 self.draft_tokens_np = async_copy_to_np(draft_tokens)
38 + # draft_tokens is a temporary allocation on the main stream and read here on
39 + # copy_stream; without record_stream, the caching allocator may reuse its
40 + # memory before the async copy executes.
41 + draft_tokens.record_stream(self.copy_stream)
42 self.copy_event.record()
43
44 def get_draft_tokens(self) -> DraftTokenIds | None:
vllm/v1/worker/gpu/warmup.py
+13 -12
@@ -30,17 +30,18 @@ def warmup_kernels(
30 pipeline parallel coordination.
31
32 The first iteration simulates a prefill with requests of
33 - 2 + num_spec_steps prompt tokens each. The second iteration simulates
34 - a decode step with all requests generating 1 + num_spec_steps tokens.
33 + decode_query_len + 1 prompt tokens each. The second iteration simulates
34 + a decode step with all requests generating decode_query_len tokens.
35 """
36 num_spec_steps = model_runner.num_speculative_steps
37 - # Use 1 + num_spec_steps + 1 tokens so the prefill batch's per-request
38 - # query length exceeds decode_query_len (= 1 + num_spec_steps), preventing
39 - # it from being misclassified as a uniform decode batch.
40 - prompt_len = 2 + num_spec_steps
37 + decode_query_len = model_runner.decode_query_len
38 + # Use decode_query_len + 1 tokens so the prefill batch's per-request query
39 + # length exceeds decode_query_len, preventing it from being misclassified as
40 + # a uniform decode batch.
41 + prompt_len = decode_query_len + 1
42 prompt_token_ids = list(range(prompt_len))
42 - # After prefill, decode generates 1 verified + num_spec_steps draft tokens.
43 - decode_len = prompt_len + 1 + num_spec_steps
43 + # After prefill, decode generates decode_query_len tokens.
44 + decode_len = prompt_len + decode_query_len
45
46 kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
47 num_kv_cache_groups = len(kv_cache_groups)
@@ -57,7 +58,7 @@ def warmup_kernels(
58 num_reqs = min(
59 model_runner.scheduler_config.max_num_seqs,
60 model_runner.scheduler_config.max_num_batched_tokens
60 - // max(prompt_len, 1 + num_spec_steps),
61 + // max(prompt_len, decode_query_len),
62 # Reserve block 0 (null block) and ensure we have enough blocks.
63 max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req),
64 )
@@ -79,7 +80,7 @@ def warmup_kernels(
80 nonlocal next_block_id
81 return list(range(next_block_id, next_block_id := next_block_id + num_blocks))
82
82 - # Step 1: Prefill all requests with 2 + num_spec_steps prompt tokens each.
83 + # Step 1: Prefill all requests with 1 + decode_query_len prompt tokens each.
84 new_reqs = [
85 NewRequestData.from_request(
86 Request(req_ids[i], prompt_token_ids, sampling_params, pooling_params),
@@ -117,7 +118,7 @@ def warmup_kernels(
118
119 worker_sample_tokens(grammar_output)
120
120 - # Step 2: Decode all requests with 1 + num_spec_steps tokens each.
121 + # Step 2: Decode all requests with decode_query_len tokens each.
122 cached_req_data = CachedRequestData.make_empty()
123 cached_req_data.req_ids = list(req_ids)
124 cached_req_data.num_computed_tokens = [prompt_len] * num_reqs
@@ -131,7 +132,7 @@ def warmup_kernels(
132 decode_output = SchedulerOutput.make_empty()
133 decode_output.scheduled_cached_reqs = cached_req_data
134 decode_output.num_scheduled_tokens = {
134 - req_id: 1 + num_spec_steps for req_id in req_ids
135 + req_id: decode_query_len for req_id in req_ids
136 }
137 if num_spec_steps > 0:
138 decode_output.scheduled_spec_decode_tokens = {
vllm/vllm_flash_attn/flash_attn_interface.py
+2
@@ -209,6 +209,7 @@ def flash_attn_varlen_func(
209 # FA4 only
210 mask_mod=None,
211 aux_tensors=None,
212 + dynamic_causal: "torch.Tensor | None" = None,
213 ):
214 """dropout_p should be set to 0.0 during evaluation
215 Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads
@@ -392,6 +393,7 @@ def flash_attn_varlen_func(
393 page_table=block_table,
394 softmax_scale=softmax_scale,
395 causal=causal,
396 + dynamic_causal=dynamic_causal,
397 softcap=softcap,
398 window_size_left=real_window_size[0] if real_window_size[0] >= 0 else None,
399 window_size_right=real_window_size[1] if real_window_size[1] >= 0 else None,