A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 02:45 UTC
[Core] NGram GPU Implementation compatible with Async Scheduler (#29184)
PatchyTIS committed
Mar 8, 2026 at 05:51 UTC
a6be75dbd2a8dd1886da725727ee178f42e3f84f
9 files changed
+940
-12
tests/v1/e2e/test_async_scheduling.py
+40
-3
@@ -98,7 +98,7 @@ def test_without_spec_decoding(
98
99
@single_gpu_only
100
@large_gpu_mark(min_gb=16)
101
-def test_with_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch):
101
+def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch):
102
"""Test consistency and acceptance rates with some different combos of
103
preemption, executor, async scheduling, prefill chunking,
104
spec decoding model length.
@@ -154,6 +154,42 @@ def test_with_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch)
154
)
155
156
157
+def test_with_ngram_gpu_spec_decoding(monkeypatch: pytest.MonkeyPatch):
158
+ """Test ngram_gpu speculative decoding with different configurations.
159
+
160
+ This test specifically validates ngram_gpu behavior with various:
161
+ - Number of speculative tokens (2-6)
162
+ - Prompt lookup window sizes (min/max)
163
+ - Async scheduling enabled (as in production)
164
+ - Different executors and chunking settings
165
+ """
166
+
167
+ # Variant with larger speculation window
168
+ ngram_gpu_config = {
169
+ "method": "ngram_gpu",
170
+ "num_speculative_tokens": 3,
171
+ "prompt_lookup_max": 3,
172
+ "prompt_lookup_min": 2,
173
+ }
174
+
175
+ # Test configurations covering various scenarios
176
+ # test_preemption, executor, async_scheduling,
177
+ # spec_config, test_prefill_chunking
178
+ test_configs = [
179
+ (False, "mp", False, None, False),
180
+ (False, "mp", False, ngram_gpu_config, False),
181
+ (True, "mp", False, ngram_gpu_config, True),
182
+ (False, "mp", True, ngram_gpu_config, False),
183
+ (True, "mp", True, ngram_gpu_config, False),
184
+ (True, "uni", True, ngram_gpu_config, False),
185
+ (True, "mp", True, ngram_gpu_config, True),
186
+ ]
187
+
188
+ # Use MODEL (Qwen) for ngram_gpu tests as it's lighter weight
189
+ # and ngram_gpu doesn't require a specific draft model
190
+ run_tests(monkeypatch, MODEL, test_configs, [{}])
191
+
192
+
193
@dynamo_config.patch(cache_size_limit=16)
194
def run_tests(
195
monkeypatch: pytest.MonkeyPatch,
@@ -282,11 +318,12 @@ def run_test(
318
else dict(gpu_memory_utilization=0.9)
319
)
320
spec_mml = (spec_config or {}).get("max_model_len")
321
+ spec_method = (spec_config or {}).get("method", "none")
322
test_config = (
323
f"executor={executor}, preemption={test_preemption}, "
324
f"async_sched={async_scheduling}, "
325
f"chunk_prefill={test_prefill_chunking}, "
289
- f"spec_decoding={spec_decoding}, spec_mml={spec_mml}"
326
+ f"spec_decoding={spec_decoding}, spec_method={spec_method}, spec_mml={spec_mml}"
327
)
328
print("-" * 80)
329
print(f"---- TESTING {test_str}: {test_config}")
@@ -294,7 +331,7 @@ def run_test(
331
332
with VllmRunner(
333
model,
297
- max_model_len=512,
334
+ max_model_len=4096,
335
enable_chunked_prefill=test_prefill_chunking,
336
# Force prefill chunking
337
max_num_batched_tokens=48 if test_prefill_chunking else None,
tests/v1/e2e/test_spec_decode.py
+28
@@ -183,6 +183,34 @@ def test_ngram_and_suffix_correctness(
183
cleanup_dist_env_and_memory()
184
185
186
+@pytest.mark.parametrize("async_scheduling", [True], ids=["async"])
187
+@single_gpu_only
188
+@large_gpu_mark(min_gb=20)
189
+def test_ngram_gpu_default_with_async_scheduling(
190
+ async_scheduling: bool,
191
+):
192
+ """
193
+ Test ngram_gpu speculative decoding (k=3) correctness with and without
194
+ async scheduling, validated via GSM8K accuracy.
195
+ Uses Qwen/Qwen3-8B (ref GSM8K accuracy: 87%-92%).
196
+ """
197
+ qwen3_model = "Qwen/Qwen3-8B"
198
+ spec_llm = LLM(
199
+ model=qwen3_model,
200
+ speculative_config={
201
+ "method": "ngram_gpu",
202
+ "prompt_lookup_max": 3,
203
+ "prompt_lookup_min": 2,
204
+ "num_speculative_tokens": 2,
205
+ },
206
+ max_model_len=4096,
207
+ async_scheduling=async_scheduling,
208
+ )
209
+ evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8)
210
+ del spec_llm
211
+ cleanup_dist_env_and_memory()
212
+
213
+
214
@single_gpu_only
215
@large_gpu_mark(min_gb=20)
216
def test_suffix_decoding_acceptance(
vllm/compilation/backends.py
+7
@@ -907,6 +907,13 @@ class VllmBackend:
907
# Honors opt-outs such as CompilationMode.NONE or VLLM_DISABLE_COMPILE_CACHE.
908
disable_cache = not is_compile_cache_enabled(self.inductor_config)
909
910
+ # TODO(patchy): ngram gpu kernel will cause vllm torch compile cache errors.
911
+ is_ngram_gpu_enabled = (
912
+ vllm_config.speculative_config is not None
913
+ and vllm_config.speculative_config.use_ngram_gpu()
914
+ )
915
+ disable_cache = disable_cache or is_ngram_gpu_enabled
916
+
917
if disable_cache:
918
logger.info_once("vLLM's torch.compile cache is disabled.", scope="local")
919
else:
vllm/config/speculative.py
+9
-1
@@ -47,6 +47,7 @@ MTPModelTypes = Literal[
47
"step3p5_mtp",
48
]
49
EagleModelTypes = Literal["eagle", "eagle3", "extract_hidden_states", MTPModelTypes]
50
+NgramGPUTypes = Literal["ngram_gpu"]
51
SpeculativeMethod = Literal[
52
"ngram",
53
"medusa",
@@ -54,6 +55,7 @@ SpeculativeMethod = Literal[
55
"draft_model",
56
"suffix",
57
EagleModelTypes,
58
+ NgramGPUTypes,
59
]
60
61
@@ -364,6 +366,8 @@ class SpeculativeConfig:
366
self.quantization = self.target_model_config.quantization
367
elif self.method in ("ngram", "[ngram]"):
368
self.model = "ngram"
369
+ elif self.method == "ngram_gpu":
370
+ self.model = "ngram_gpu"
371
elif self.method == "suffix":
372
self.model = "suffix"
373
elif self.method == "extract_hidden_states":
@@ -374,8 +378,9 @@ class SpeculativeConfig:
378
)
379
380
if self.method in ("ngram", "[ngram]"):
377
- # Unified to "ngram" internally
381
self.method = "ngram"
382
+
383
+ if self.method in ("ngram", "ngram_gpu"):
384
# Set default values if not provided
385
if self.prompt_lookup_min is None and self.prompt_lookup_max is None:
386
# TODO(woosuk): Tune these values. They are arbitrarily chosen.
@@ -832,6 +837,9 @@ class SpeculativeConfig:
837
def uses_extract_hidden_states(self) -> bool:
838
return self.method == "extract_hidden_states"
839
840
+ def use_ngram_gpu(self) -> bool:
841
+ return self.method == "ngram_gpu"
842
+
843
def __repr__(self) -> str:
844
method = self.method
845
model = (
vllm/config/vllm.py
+5
-2
@@ -41,7 +41,7 @@ from .offload import OffloadConfig
41
from .parallel import ParallelConfig
42
from .profiler import ProfilerConfig
43
from .scheduler import SchedulerConfig
44
-from .speculative import EagleModelTypes, SpeculativeConfig
44
+from .speculative import EagleModelTypes, NgramGPUTypes, SpeculativeConfig
45
from .structured_outputs import StructuredOutputsConfig
46
from .utils import SupportsHash, config, replace
47
from .weight_transfer import WeightTransferConfig
@@ -696,11 +696,13 @@ class VllmConfig:
696
if self.speculative_config is not None:
697
if (
698
self.speculative_config.method not in get_args(EagleModelTypes)
699
+ and self.speculative_config.method not in get_args(NgramGPUTypes)
700
and self.speculative_config.method != "draft_model"
701
):
702
raise ValueError(
703
"Currently, async scheduling is only supported "
703
- "with EAGLE/MTP/Draft Model kind of speculative decoding."
704
+ "with EAGLE/MTP/Draft Model/NGram GPU kind of "
705
+ "speculative decoding"
706
)
707
if self.speculative_config.disable_padded_drafter_batch:
708
raise ValueError(
@@ -718,6 +720,7 @@ class VllmConfig:
720
if (
721
self.speculative_config is not None
722
and self.speculative_config.method not in get_args(EagleModelTypes)
723
+ and self.speculative_config.method not in get_args(NgramGPUTypes)
724
):
725
logger.warning_once(
726
"Async scheduling not supported with %s-based "
vllm/tool_parsers/hermes_tool_parser.py
+2
@@ -385,6 +385,7 @@ class Hermes2ProToolParser(ToolParser):
385
prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get(
386
"arguments"
387
)
388
+ assert current_tool_call is not None
389
cur_arguments = current_tool_call.get("arguments")
390
391
logger.debug("diffing old arguments: %s", prev_arguments)
@@ -489,6 +490,7 @@ class Hermes2ProToolParser(ToolParser):
490
491
# handle saving the state for the current tool into
492
# the "prev" list for use in diffing for the next iteration
493
+ assert isinstance(current_tool_call, dict)
494
if self.current_tool_id == len(self.prev_tool_call_arr) - 1:
495
self.prev_tool_call_arr[self.current_tool_id] = current_tool_call
496
else:
vllm/v1/spec_decode/ngram_proposer_gpu.py
new
+660
@@ -0,0 +1,660 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+GPU-accelerated N-gram proposer using fully async PyTorch tensor operations.
5
+
6
+This version uses a fully vectorized approach with unfold and argmax for
7
+finding the first match across all sequences in parallel.
8
+"""
9
+
10
+import torch
11
+from torch import nn
12
+
13
+from vllm.compilation.decorators import support_torch_compile
14
+from vllm.config import (
15
+ CompilationConfig,
16
+ CompilationMode,
17
+ CUDAGraphMode,
18
+ VllmConfig,
19
+)
20
+from vllm.forward_context import set_forward_context
21
+from vllm.v1.core.sched.output import SchedulerOutput
22
+from vllm.v1.utils import record_function_or_nullcontext
23
+from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
24
+
25
+
26
+@support_torch_compile()
27
+class NgramGPUKernel(nn.Module):
28
+ """GPU-accelerated N-gram proposer using fully async tensor operations."""
29
+
30
+ def __init__(
31
+ self, vllm_config: VllmConfig, prefix: str = "", device: torch.device = "cuda"
32
+ ):
33
+ super().__init__()
34
+
35
+ assert vllm_config.speculative_config is not None
36
+ assert vllm_config.speculative_config.prompt_lookup_min is not None
37
+ assert vllm_config.speculative_config.prompt_lookup_max is not None
38
+
39
+ self.min_n = vllm_config.speculative_config.prompt_lookup_min
40
+ self.max_n = vllm_config.speculative_config.prompt_lookup_max
41
+ self.k = vllm_config.speculative_config.num_speculative_tokens
42
+ self.max_model_len = vllm_config.model_config.max_model_len
43
+ self.max_num_seqs = vllm_config.scheduler_config.max_num_seqs
44
+ self.device = device
45
+
46
+ def _find_first_and_extract_all_n_parallel(
47
+ self,
48
+ token_ids: torch.Tensor,
49
+ seq_lengths: torch.Tensor,
50
+ min_ngram_len: int,
51
+ max_ngram_len: int,
52
+ num_draft_tokens: int,
53
+ ) -> torch.Tensor:
54
+ """
55
+ Find suffix n-gram matches and extract following tokens.
56
+ Searches for the earliest prior occurrence of the trailing n-gram,
57
+ tries multiple lengths, and picks the longest valid match.
58
+
59
+ Args:
60
+ token_ids: Token IDs for each sequence
61
+ seq_lengths: Actual length of each sequence (excluding padding)
62
+ min_ngram_len: Minimum n-gram size to search for (e.g., 2)
63
+ max_ngram_len: Maximum n-gram size to search for (e.g., 5)
64
+ num_draft_tokens: Number of tokens to extract after match (k)
65
+
66
+ Returns:
67
+ Draft token predictions; -1 means invalid/no match.
68
+ """
69
+ batch_size = token_ids.shape[0]
70
+ max_seq_len = token_ids.shape[1]
71
+ device = token_ids.device
72
+ num_ngram_sizes = max_ngram_len - min_ngram_len + 1
73
+
74
+ # All n-gram sizes to try.
75
+ ngram_lengths = torch.arange(min_ngram_len, max_ngram_len + 1, device=device)
76
+ batch_indices = torch.arange(batch_size, device=device)
77
+
78
+ # Earliest match per (sequence, ngram_len); -1 means no match.
79
+ first_match_positions = torch.full(
80
+ (batch_size, num_ngram_sizes), -1, dtype=torch.long, device=device
81
+ )
82
+
83
+ for i, ngram_len in enumerate(range(min_ngram_len, max_ngram_len + 1)):
84
+ # Sliding windows of size ngram_len; unfold is O(1) view.
85
+ search_windows = token_ids.unfold(1, ngram_len, 1)
86
+ num_windows = search_windows.shape[1]
87
+
88
+ # Trailing suffix (last ngram_len tokens) for each sequence.
89
+ suffix_starts = seq_lengths - ngram_len
90
+ suffix_indices = suffix_starts.unsqueeze(1) + torch.arange(
91
+ ngram_len, device=device
92
+ )
93
+ suffix = torch.gather(token_ids, 1, suffix_indices.clamp(min=0))
94
+
95
+ # Window matches for each sequence.
96
+ matches = (search_windows == suffix.unsqueeze(1)).all(dim=-1)
97
+
98
+ # Match must leave room for at least one draft token.
99
+ max_valid_suffix_start = seq_lengths - ngram_len - 1
100
+ window_positions = torch.arange(num_windows, device=device)
101
+ valid_mask = window_positions <= max_valid_suffix_start.unsqueeze(1)
102
+ final_matches = matches & valid_mask
103
+
104
+ # Find earliest match (argmax=0 when empty; verify with has_match).
105
+ first_match_idx = torch.argmax(final_matches.int(), dim=1)
106
+ has_match = final_matches[batch_indices, first_match_idx]
107
+
108
+ # Store valid match positions (window index = position).
109
+ first_match_positions[:, i] = torch.where(has_match, first_match_idx, -1)
110
+
111
+ # Select the longest n-gram with a match.
112
+ best_ngram_idx = (first_match_positions >= 0).int().flip(dims=[1]).argmax(dim=1)
113
+ best_ngram_idx = num_ngram_sizes - 1 - best_ngram_idx # Flip back
114
+
115
+ # Match position for the best n-gram.
116
+ best_match_pos = first_match_positions[batch_indices, best_ngram_idx]
117
+
118
+ # Avoid data-dependent branching.
119
+ has_any_match = best_match_pos >= 0
120
+
121
+ # Length of the best matching n-gram.
122
+ best_ngram_lengths = ngram_lengths[best_ngram_idx]
123
+
124
+ # Start position right after the matched suffix.
125
+ draft_start = torch.where(
126
+ has_any_match,
127
+ best_match_pos + best_ngram_lengths,
128
+ torch.zeros_like(best_match_pos),
129
+ )
130
+ tokens_available = seq_lengths - draft_start
131
+
132
+ # Gather indices for draft tokens.
133
+ draft_indices = draft_start.unsqueeze(1) + torch.arange(
134
+ num_draft_tokens, device=device
135
+ )
136
+ draft_indices = draft_indices.clamp(min=0, max=max_seq_len - 1)
137
+
138
+ # Extract draft tokens; gather always runs.
139
+ draft_tokens = torch.gather(token_ids, 1, draft_indices)
140
+
141
+ # Mask positions beyond available tokens.
142
+ position_indices = torch.arange(num_draft_tokens, device=device).unsqueeze(0)
143
+ valid_positions = position_indices < tokens_available.unsqueeze(1)
144
+
145
+ draft_tokens = torch.where(
146
+ valid_positions,
147
+ draft_tokens,
148
+ torch.full_like(draft_tokens, -1),
149
+ )
150
+
151
+ # If no match, mask all positions.
152
+ draft_tokens = torch.where(
153
+ has_any_match.unsqueeze(1),
154
+ draft_tokens,
155
+ torch.full_like(draft_tokens, -1),
156
+ )
157
+
158
+ return draft_tokens
159
+
160
+ def forward(
161
+ self,
162
+ num_tokens_no_spec: torch.Tensor,
163
+ token_ids_gpu: torch.Tensor,
164
+ combined_mask: torch.Tensor,
165
+ ) -> tuple[torch.Tensor, torch.Tensor]:
166
+ """
167
+ Forward pass for N-gram proposal using GPU tensor operations.
168
+
169
+ Args:
170
+ num_tokens_no_spec: Number of tokens for each sequence [batch_size]
171
+ token_ids_gpu: Token IDs [batch_size, max_len]
172
+ combined_mask: Whether each sequence is valid for spec decode [batch_size]
173
+
174
+ Returns:
175
+ draft_tokens: [batch_size, k] on GPU
176
+ num_valid_draft_tokens: [batch_size] int32 on GPU, count of
177
+ leading valid (non -1) tokens per request.
178
+ """
179
+
180
+ device = token_ids_gpu.device
181
+
182
+ # Infer batch size to preserve dynamic shape.
183
+ actual_batch_size = token_ids_gpu.shape[0]
184
+
185
+ # Allocate in forward so torch.compile can optimize.
186
+ # NOTE(patchy): Do NOT pre-allocate this as a buffer
187
+ # it breaks torch.compile
188
+ draft_tokens = torch.full(
189
+ (actual_batch_size, self.k), -1, dtype=torch.int32, device=device
190
+ )
191
+
192
+ results = self._find_first_and_extract_all_n_parallel(
193
+ token_ids_gpu,
194
+ num_tokens_no_spec,
195
+ min_ngram_len=self.min_n,
196
+ max_ngram_len=self.max_n,
197
+ num_draft_tokens=self.k,
198
+ )
199
+
200
+ draft_tokens = torch.where(combined_mask.unsqueeze(1), results, -1)
201
+
202
+ # Count leading contiguous valid (non -1) tokens per request.
203
+ is_valid = draft_tokens != -1 # [batch, k]
204
+ cum_valid = is_valid.int().cumsum(dim=1) # [batch, k]
205
+ positions = torch.arange(1, self.k + 1, device=device).unsqueeze(0)
206
+ num_valid_draft_tokens = (cum_valid == positions).int().sum(dim=1)
207
+
208
+ return draft_tokens, num_valid_draft_tokens
209
+
210
+ def load_model(self, *args, **kwargs):
211
+ """No model to load for N-gram proposer."""
212
+ pass
213
+
214
+
215
+class NgramProposerGPU:
216
+ def __init__(self, vllm_config: VllmConfig, device: torch.device, runner=None):
217
+ assert vllm_config.speculative_config is not None
218
+ assert vllm_config.speculative_config.prompt_lookup_min is not None
219
+ assert vllm_config.speculative_config.prompt_lookup_max is not None
220
+
221
+ compilation_config = CompilationConfig(
222
+ mode=CompilationMode.VLLM_COMPILE,
223
+ custom_ops=["none"],
224
+ splitting_ops=[],
225
+ compile_sizes=[],
226
+ inductor_compile_config={
227
+ "enable_auto_functionalized_v2": False,
228
+ "max_autotune": True,
229
+ "aggressive_fusion": True,
230
+ "triton.autotune_pointwise": True,
231
+ "coordinate_descent_tuning": True,
232
+ "use_mixed_mm": False,
233
+ },
234
+ cudagraph_mode=CUDAGraphMode.NONE,
235
+ )
236
+ model_config = vllm_config.model_config
237
+ speculative_config = vllm_config.speculative_config
238
+ scheduler_config = vllm_config.scheduler_config
239
+
240
+ self.vllm_config = VllmConfig(
241
+ compilation_config=compilation_config,
242
+ model_config=model_config,
243
+ speculative_config=speculative_config,
244
+ scheduler_config=scheduler_config,
245
+ )
246
+
247
+ self.min_n = vllm_config.speculative_config.prompt_lookup_min
248
+ self.max_n = vllm_config.speculative_config.prompt_lookup_max
249
+ self.k = vllm_config.speculative_config.num_speculative_tokens
250
+ self.max_model_len = vllm_config.model_config.max_model_len
251
+ self.max_num_seqs = vllm_config.scheduler_config.max_num_seqs
252
+ self.device = device
253
+
254
+ self.kernel = NgramGPUKernel(
255
+ vllm_config=self.vllm_config, prefix="ngram_gpu_kernel", device=device
256
+ )
257
+ self.kernel.to(device)
258
+ self.kernel.eval()
259
+
260
+ self._dummy_run()
261
+
262
+ def _dummy_run(self):
263
+ token_ids, num_tokens, sampled_flags, valid_mask = self._generate_dummy_data(
264
+ batch_size=self.max_num_seqs,
265
+ max_seq_len=self.max_model_len,
266
+ pattern_len=self.k,
267
+ device=self.device,
268
+ )
269
+
270
+ combined_mask = sampled_flags & valid_mask & (num_tokens >= self.min_n)
271
+
272
+ for _ in range(3):
273
+ with set_forward_context(None, self.vllm_config):
274
+ _, _ = self.kernel(num_tokens, token_ids, combined_mask)
275
+
276
+ def _generate_dummy_data(
277
+ self,
278
+ batch_size: int,
279
+ max_seq_len: int,
280
+ pattern_len: int,
281
+ device: str = "cuda",
282
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
283
+ """
284
+ Generate random test data with n-gram repetitions.
285
+
286
+ Args:
287
+ batch_size: Number of sequences in the batch
288
+ max_seq_len: Maximum sequence length
289
+ pattern_len: Length of patterns to inject for matching
290
+ device: Device to place tensors on
291
+
292
+ Returns:
293
+ token_ids: [batch_size, max_seq_len] tensor
294
+ num_tokens: [batch_size] tensor
295
+ sampled_flags: [batch_size] bool tensor
296
+ valid_mask: [batch_size] bool tensor
297
+ """
298
+ token_ids = torch.zeros(
299
+ batch_size,
300
+ max_seq_len,
301
+ dtype=torch.int32,
302
+ device=device,
303
+ )
304
+
305
+ num_tokens = torch.randint(
306
+ pattern_len, max_seq_len, (batch_size,), dtype=torch.int32, device=device
307
+ )
308
+
309
+ sampled_flags = torch.ones(batch_size, dtype=torch.bool, device=device)
310
+ valid_mask = torch.ones(batch_size, dtype=torch.bool, device=device)
311
+
312
+ return token_ids, num_tokens, sampled_flags, valid_mask
313
+
314
+ def propose(
315
+ self,
316
+ num_tokens_no_spec: torch.Tensor, # [batch_size]
317
+ token_ids_gpu: torch.Tensor, # [batch_size, max_len]
318
+ valid_sampled_token_ids_gpu: torch.Tensor, # [batch_size, num_spec_tokens + 1]
319
+ valid_sampled_tokens_count: torch.Tensor, # [batch_size]
320
+ ) -> tuple[torch.Tensor, torch.Tensor]:
321
+ """
322
+ Propose draft tokens using GPU-accelerated n-gram matching.
323
+
324
+ Scatter sampled tokens into `token_ids_gpu`, compute temporary
325
+ updated lengths, then run the kernel.
326
+
327
+ Args:
328
+ num_tokens_no_spec: Number of tokens per sequence (read-only)
329
+ token_ids_gpu: Token IDs tensor (modified in-place with new tokens)
330
+ valid_sampled_token_ids_gpu: Newly sampled tokens to scatter
331
+ valid_sampled_tokens_count: Count of valid tokens per sequence
332
+
333
+ Returns:
334
+ draft_tokens: Proposed draft token IDs [batch_size, k]
335
+ num_valid_draft_tokens: Count of leading valid draft tokens
336
+ per request [batch_size]
337
+ """
338
+ assert token_ids_gpu.device == self.device
339
+ assert num_tokens_no_spec.device == self.device
340
+
341
+ batch_size = num_tokens_no_spec.shape[0]
342
+ max_seq_len = token_ids_gpu.shape[1]
343
+ max_new_tokens = valid_sampled_token_ids_gpu.shape[1] # num_spec_tokens + 1
344
+
345
+ # Scatter newly sampled tokens into token_ids_gpu.
346
+ offsets = torch.arange(max_new_tokens, device=self.device)
347
+ write_positions = num_tokens_no_spec.unsqueeze(1) + offsets.unsqueeze(0)
348
+ valid_write_mask = offsets.unsqueeze(0) < valid_sampled_tokens_count.unsqueeze(
349
+ 1
350
+ )
351
+ in_bounds = write_positions < max_seq_len
352
+ scatter_mask = (
353
+ valid_write_mask & (valid_sampled_token_ids_gpu != -1) & in_bounds
354
+ )
355
+
356
+ write_positions_long = write_positions.clamp(max=max_seq_len - 1).long()
357
+ existing_values = token_ids_gpu.gather(1, write_positions_long)
358
+
359
+ tokens_cast = valid_sampled_token_ids_gpu.to(token_ids_gpu.dtype)
360
+ tokens_to_scatter = torch.where(
361
+ scatter_mask,
362
+ tokens_cast,
363
+ existing_values,
364
+ )
365
+ token_ids_gpu.scatter_(1, write_positions_long, tokens_to_scatter)
366
+
367
+ num_tokens_tmp = num_tokens_no_spec + valid_sampled_tokens_count
368
+
369
+ # Compute validity masks.
370
+ sampled_flags = valid_sampled_tokens_count > 0
371
+ valid_mask = torch.ones(batch_size, dtype=torch.bool, device=self.device)
372
+
373
+ with set_forward_context(None, self.vllm_config):
374
+ combined_mask = sampled_flags & valid_mask & (num_tokens_tmp >= self.min_n)
375
+
376
+ with record_function_or_nullcontext("ngram_proposer_gpu: kernel"):
377
+ draft_tokens, num_valid_draft_tokens = self.kernel(
378
+ num_tokens_tmp,
379
+ token_ids_gpu,
380
+ combined_mask,
381
+ )
382
+
383
+ return draft_tokens, num_valid_draft_tokens
384
+
385
+ def update_token_ids_ngram(
386
+ self,
387
+ sampled_token_ids: torch.Tensor | list[list[int]],
388
+ gpu_input_batch: InputBatch,
389
+ token_ids_gpu: torch.Tensor,
390
+ num_tokens_no_spec: torch.Tensor,
391
+ discard_request_mask: torch.Tensor,
392
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
393
+ """
394
+ Prepare speculative decoding inputs on device:
395
+ compute next token ids and valid counts, honoring discarded requests
396
+ and rejected tokens, without CPU-GPU sync.
397
+ """
398
+ num_reqs = gpu_input_batch.num_reqs
399
+
400
+ if isinstance(sampled_token_ids, list):
401
+ # When disable_padded_drafter_batch=True, sampled_token_ids is
402
+ # an irregular list[list[int]] where sublists may have different
403
+ # lengths (including empty lists for discarded requests).
404
+ # Pad all sublists to the same length with -1 before converting
405
+ # to tensor.
406
+ max_len = max(
407
+ (len(sublist) for sublist in sampled_token_ids),
408
+ default=0,
409
+ )
410
+ # Ensure at least length 1 for tensor creation
411
+ max_len = max(max_len, 1)
412
+ padded_list = [
413
+ sublist + [-1] * (max_len - len(sublist))
414
+ for sublist in sampled_token_ids
415
+ ]
416
+ sampled_token_ids = torch.tensor(
417
+ padded_list, dtype=torch.int32, device=self.device
418
+ )
419
+ assert isinstance(sampled_token_ids, torch.Tensor), (
420
+ "sampled_token_ids should be a torch.Tensor for ngram_gpu"
421
+ )
422
+
423
+ # Backup last valid token before speculative tokens.
424
+ backup_indices = (num_tokens_no_spec[:num_reqs] - 1).clamp(min=0).long()
425
+ backup_next_token_ids = torch.gather(
426
+ token_ids_gpu[:num_reqs], dim=1, index=backup_indices.unsqueeze(1)
427
+ ).squeeze(1)
428
+
429
+ valid_sampled_token_ids_gpu = sampled_token_ids.clone()
430
+ # Invalidate sampled tokens for discarded requests.
431
+ discard_mask_expanded = discard_request_mask[:num_reqs].unsqueeze(1)
432
+ valid_sampled_token_ids_gpu.masked_fill_(discard_mask_expanded, -1)
433
+
434
+ # Mask valid tokens within each request.
435
+ valid_mask = (valid_sampled_token_ids_gpu != -1) & (
436
+ valid_sampled_token_ids_gpu < gpu_input_batch.vocab_size
437
+ )
438
+
439
+ # Count valid tokens per request.
440
+ valid_sampled_tokens_count = valid_mask.sum(dim=1)
441
+
442
+ # Rightmost valid index per row.
443
+ last_valid_indices = valid_sampled_tokens_count - 1
444
+ last_valid_indices_safe = torch.clamp(last_valid_indices, min=0)
445
+
446
+ # Last valid token from each row; undefined if none.
447
+ selected_tokens = torch.gather(
448
+ valid_sampled_token_ids_gpu, 1, last_valid_indices_safe.unsqueeze(1)
449
+ ).squeeze(1)
450
+
451
+ # Use last token if valid; otherwise fallback to backup.
452
+ next_token_ids = torch.where(
453
+ last_valid_indices != -1,
454
+ selected_tokens,
455
+ backup_next_token_ids,
456
+ )
457
+
458
+ return next_token_ids, valid_sampled_tokens_count, valid_sampled_token_ids_gpu
459
+
460
+ def load_model(self, *args, **kwargs):
461
+ self.kernel.load_model(*args, **kwargs)
462
+
463
+
464
+def update_scheduler_for_invalid_drafts(
465
+ num_valid_draft_tokens_event: torch.cuda.Event,
466
+ num_valid_draft_tokens_cpu: torch.Tensor,
467
+ scheduler_output: "SchedulerOutput",
468
+ req_id_to_index: dict[str, int],
469
+) -> None:
470
+ """Trim invalid speculative slots using per-request valid draft counts.
471
+
472
+ Args:
473
+ num_valid_draft_tokens_event: Event for async D2H completion.
474
+ num_valid_draft_tokens_cpu: CPU buffer of valid draft counts.
475
+ scheduler_output: Scheduler metadata to update in-place.
476
+ req_id_to_index: Request-id to batch-index mapping.
477
+ """
478
+ req_data = scheduler_output.scheduled_cached_reqs
479
+ num_valid_draft_tokens_event.synchronize()
480
+
481
+ for req_id in req_data.req_ids:
482
+ req_index = req_id_to_index.get(req_id)
483
+ if req_index is None:
484
+ continue
485
+
486
+ spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
487
+ if spec_token_ids is None:
488
+ continue
489
+
490
+ scheduled_k = len(spec_token_ids)
491
+
492
+ valid_k = int(num_valid_draft_tokens_cpu[req_index].item())
493
+ valid_k = max(0, min(valid_k, scheduled_k))
494
+
495
+ tokens_to_trim = scheduled_k - valid_k
496
+ scheduler_output.total_num_scheduled_tokens -= tokens_to_trim
497
+ scheduler_output.num_scheduled_tokens[req_id] -= tokens_to_trim
498
+
499
+ if valid_k == 0:
500
+ scheduler_output.scheduled_spec_decode_tokens.pop(req_id, None)
501
+ else:
502
+ scheduler_output.scheduled_spec_decode_tokens[req_id] = spec_token_ids[
503
+ :valid_k
504
+ ]
505
+
506
+
507
+def update_ngram_gpu_tensors_incremental(
508
+ input_batch: InputBatch,
509
+ token_ids_gpu_tensor: torch.Tensor,
510
+ num_tokens_no_spec_gpu: torch.Tensor,
511
+ new_reqs: list[CachedRequestState],
512
+ device: torch.device,
513
+ _pinned_idx_buf: torch.Tensor,
514
+ _pinned_val_buf: torch.Tensor,
515
+) -> None:
516
+ """Incrementally update token_ids_gpu_tensor and num_tokens_no_spec_gpu
517
+ for ngram GPU proposer.
518
+ """
519
+ prev_req_id_to_index = input_batch.prev_req_id_to_index
520
+ curr_req_id_to_index = input_batch.req_id_to_index
521
+
522
+ if not curr_req_id_to_index:
523
+ return
524
+
525
+ active_indices = list(curr_req_id_to_index.values())
526
+ n_active = len(active_indices)
527
+
528
+ # Use resident pinned buffers to avoid per-call allocation.
529
+ active_idx_cpu = _pinned_idx_buf[:n_active]
530
+ active_idx_cpu.copy_(torch.as_tensor(active_indices, dtype=torch.long))
531
+
532
+ active_idx_gpu = active_idx_cpu.to(device=device, non_blocking=True)
533
+
534
+ new_req_ids = {req.req_id for req in new_reqs}
535
+
536
+ # First run, no previous state.
537
+ if prev_req_id_to_index is None:
538
+ for idx in active_indices:
539
+ num_tokens = input_batch.num_tokens_no_spec[idx]
540
+ if num_tokens > 0:
541
+ token_ids_gpu_tensor[idx, :num_tokens].copy_(
542
+ input_batch.token_ids_cpu_tensor[idx, :num_tokens],
543
+ non_blocking=True,
544
+ )
545
+
546
+ _sync_num_tokens(
547
+ input_batch,
548
+ num_tokens_no_spec_gpu,
549
+ active_idx_cpu,
550
+ active_idx_gpu,
551
+ n_active,
552
+ device,
553
+ _pinned_val_buf,
554
+ )
555
+ return
556
+
557
+ # Detect index changes for reorder.
558
+ reorder_src: list[int] = []
559
+ reorder_dst: list[int] = []
560
+
561
+ for req_id, curr_idx in curr_req_id_to_index.items():
562
+ if req_id in new_req_ids:
563
+ continue
564
+ prev_idx = prev_req_id_to_index.get(req_id)
565
+ if prev_idx is not None and prev_idx != curr_idx:
566
+ reorder_src.append(prev_idx)
567
+ reorder_dst.append(curr_idx)
568
+
569
+ if reorder_src:
570
+ src_tensor = torch.tensor(reorder_src, dtype=torch.long, device=device)
571
+ dst_tensor = torch.tensor(reorder_dst, dtype=torch.long, device=device)
572
+
573
+ temp_token_ids = token_ids_gpu_tensor[src_tensor].clone()
574
+ temp_num_tokens = num_tokens_no_spec_gpu[src_tensor].clone()
575
+
576
+ token_ids_gpu_tensor[dst_tensor] = temp_token_ids
577
+ num_tokens_no_spec_gpu[dst_tensor] = temp_num_tokens
578
+
579
+ # Full copy for new/resumed requests.
580
+ for req_state in new_reqs:
581
+ new_req_idx = curr_req_id_to_index.get(req_state.req_id)
582
+ if new_req_idx is None:
583
+ continue
584
+
585
+ num_tokens = input_batch.num_tokens_no_spec[new_req_idx]
586
+ if num_tokens > 0:
587
+ token_ids_gpu_tensor[new_req_idx, :num_tokens].copy_(
588
+ input_batch.token_ids_cpu_tensor[new_req_idx, :num_tokens],
589
+ non_blocking=True,
590
+ )
591
+
592
+ # Always batch-sync sequence lengths from CPU for ALL active requests.
593
+ _sync_num_tokens(
594
+ input_batch,
595
+ num_tokens_no_spec_gpu,
596
+ active_idx_cpu,
597
+ active_idx_gpu,
598
+ n_active,
599
+ device,
600
+ _pinned_val_buf,
601
+ )
602
+
603
+
604
+def _sync_num_tokens(
605
+ input_batch: InputBatch,
606
+ num_tokens_no_spec_gpu: torch.Tensor,
607
+ active_idx_cpu: torch.Tensor,
608
+ active_idx_gpu: torch.Tensor,
609
+ n_active: int,
610
+ device: torch.device,
611
+ _pinned_val_buf: torch.Tensor,
612
+) -> None:
613
+ """Batch-sync GPU sequence lengths from CPU source of truth.
614
+
615
+ Inputs:
616
+ input_batch: Batch container with CPU length tensor.
617
+ num_tokens_no_spec_gpu: Destination GPU length tensor.
618
+ active_idx_cpu: Active request indices on CPU.
619
+ active_idx_gpu: Active request indices on GPU.
620
+ n_active: Number of active requests.
621
+ device: Target CUDA device.
622
+ _pinned_val_buf: Resident pinned int32 staging buffer.
623
+ Outputs:
624
+ None (updates num_tokens_no_spec_gpu in-place).
625
+ """
626
+ src_cpu = input_batch.num_tokens_no_spec_cpu_tensor
627
+ vals = _pinned_val_buf[:n_active]
628
+ vals.copy_(src_cpu.index_select(0, active_idx_cpu))
629
+
630
+ num_tokens_no_spec_gpu.index_copy_(
631
+ 0,
632
+ active_idx_gpu,
633
+ vals.to(device=device, non_blocking=True),
634
+ )
635
+
636
+
637
+def copy_num_valid_draft_tokens(
638
+ num_valid_draft_tokens_cpu: torch.Tensor,
639
+ num_valid_draft_tokens_copy_stream: torch.cuda.Stream,
640
+ num_valid_draft_tokens_event: torch.cuda.Event,
641
+ num_valid_draft_tokens: torch.Tensor | None,
642
+ batch_size: int,
643
+) -> None:
644
+ """
645
+ Async D2H copy of per-request valid draft counts.
646
+ """
647
+ if num_valid_draft_tokens is None:
648
+ return
649
+
650
+ num_reqs_to_copy = min(batch_size, num_valid_draft_tokens.shape[0])
651
+ if num_reqs_to_copy <= 0:
652
+ return
653
+
654
+ default_stream = torch.cuda.current_stream()
655
+ with torch.cuda.stream(num_valid_draft_tokens_copy_stream):
656
+ num_valid_draft_tokens_copy_stream.wait_stream(default_stream)
657
+ num_valid_draft_tokens_cpu[:num_reqs_to_copy].copy_(
658
+ num_valid_draft_tokens[:num_reqs_to_copy], non_blocking=True
659
+ )
660
+ num_valid_draft_tokens_event.record()
vllm/v1/worker/gpu_input_batch.py
+7
-1
@@ -127,7 +127,13 @@ class InputBatch:
127
# allocation if max_model_len is big.
128
# Maps req_index -> tensor of shape (num_prompt_tokens, hidden_size)
129
self.req_prompt_embeds: dict[int, torch.Tensor] = {}
130
- self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
130
+ self.num_tokens_no_spec_cpu_tensor = torch.zeros(
131
+ (max_num_reqs,),
132
+ device="cpu",
133
+ dtype=torch.int32,
134
+ pin_memory=pin_memory,
135
+ )
136
+ self.num_tokens_no_spec = self.num_tokens_no_spec_cpu_tensor.numpy()
137
self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
138
self.num_computed_tokens_cpu_tensor = torch.zeros(
139
(max_num_reqs,),
vllm/v1/worker/gpu_model_runner.py
+182
-5
@@ -10,7 +10,7 @@ from collections import defaultdict
10
from collections.abc import Iterable, Iterator, Sequence
11
from contextlib import contextmanager
12
from copy import copy, deepcopy
13
-from dataclasses import dataclass
13
+from dataclasses import dataclass, replace
14
from functools import reduce
15
from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast
16
@@ -164,6 +164,12 @@ from vllm.v1.spec_decode.eagle import EagleProposer
164
from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
165
from vllm.v1.spec_decode.medusa import MedusaProposer
166
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
167
+from vllm.v1.spec_decode.ngram_proposer_gpu import (
168
+ NgramProposerGPU,
169
+ copy_num_valid_draft_tokens,
170
+ update_ngram_gpu_tensors_incremental,
171
+ update_scheduler_for_invalid_drafts,
172
+)
173
from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer
174
from vllm.v1.structured_output.utils import apply_grammar_bitmask
175
from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext
@@ -424,7 +430,7 @@ class GPUModelRunner(
430
431
# Broadcast PP output for external_launcher (torchrun)
432
# to make sure we are synced across pp ranks
427
- # TODO: Support overlapping mirco-batches
433
+ # TODO: Support overlapping micro-batches
434
# https://github.com/vllm-project/vllm/issues/18019
435
self.broadcast_pp_output = (
436
self.parallel_config.distributed_executor_backend == "external_launcher"
@@ -493,6 +499,7 @@ class GPUModelRunner(
499
if self.speculative_config and get_pp_group().is_last_rank:
500
self.drafter: (
501
NgramProposer # noqa: F823
502
+ | NgramProposerGPU
503
| SuffixDecodingProposer
504
| EagleProposer
505
| DraftModelProposer
@@ -509,6 +516,23 @@ class GPUModelRunner(
516
device=self.device,
517
runner=self,
518
)
519
+ elif self.speculative_config.use_ngram_gpu():
520
+ self.drafter = NgramProposerGPU(self.vllm_config, self.device, self)
521
+ self.num_tokens_no_spec_gpu = torch.zeros(
522
+ self.max_num_reqs, dtype=torch.int32, device=device
523
+ )
524
+ self.token_ids_gpu_tensor = torch.zeros(
525
+ self.max_num_reqs,
526
+ self.max_model_len,
527
+ dtype=torch.int32,
528
+ device=device,
529
+ )
530
+ self._ngram_pinned_idx_buf = torch.zeros(
531
+ self.max_num_reqs, dtype=torch.long, pin_memory=True
532
+ )
533
+ self._ngram_pinned_val_buf = torch.zeros(
534
+ self.max_num_reqs, dtype=torch.int32, pin_memory=True
535
+ )
536
elif self.speculative_config.method == "suffix":
537
self.drafter = SuffixDecodingProposer(self.vllm_config)
538
elif self.speculative_config.use_eagle():
@@ -564,7 +588,7 @@ class GPUModelRunner(
588
)
589
self.input_batch = InputBatch(
590
max_num_reqs=self.max_num_reqs,
567
- # We need to use the encoder length for encoder-decoer
591
+ # We need to use the encoder length for encoder-decoder
592
# because of KV cache for cross-attention.
593
max_model_len=max(self.max_model_len, self.max_encoder_len),
594
max_num_batched_tokens=self.max_num_tokens,
@@ -721,6 +745,21 @@ class GPUModelRunner(
745
746
# Cached outputs.
747
self._draft_token_ids: list[list[int]] | torch.Tensor | None = None
748
+ # N-gram GPU path: async D2H buffer/event for per-request valid draft counts.
749
+ self._num_valid_draft_tokens: torch.Tensor | None = None
750
+ self._num_valid_draft_tokens_cpu: torch.Tensor | None = None
751
+ self._num_valid_draft_tokens_event: torch.cuda.Event | None = None
752
+ self._num_valid_draft_tokens_copy_stream: torch.cuda.Stream | None = None
753
+ if (
754
+ self.speculative_config is not None
755
+ and self.speculative_config.use_ngram_gpu()
756
+ ):
757
+ self._num_valid_draft_tokens_cpu = torch.empty(
758
+ self.max_num_reqs, dtype=torch.int32, pin_memory=self.pin_memory
759
+ )
760
+ self._num_valid_draft_tokens_event = torch.cuda.Event()
761
+ self._num_valid_draft_tokens_copy_stream = torch.cuda.Stream()
762
+
763
self._draft_token_req_ids: list[str] | None = None
764
self.transfer_event = torch.Event()
765
self.sampled_token_ids_pinned_cpu = torch.empty(
@@ -992,6 +1031,13 @@ class GPUModelRunner(
1031
for req_id in unscheduled_req_ids:
1032
self.input_batch.remove_request(req_id)
1033
1034
+ is_ngram_gpu = (
1035
+ self.speculative_config is not None
1036
+ and self.speculative_config.use_ngram_gpu()
1037
+ )
1038
+ if is_ngram_gpu:
1039
+ ngram_gpu_new_reqs: list[CachedRequestState] = []
1040
+
1041
reqs_to_add: list[CachedRequestState] = []
1042
# Add new requests to the cached states.
1043
for new_req_data in scheduler_output.scheduled_new_reqs:
@@ -1054,12 +1100,31 @@ class GPUModelRunner(
1100
self._init_xdrope_positions(req_state)
1101
1102
reqs_to_add.append(req_state)
1103
+ # Track new requests for ngram_gpu full tensor copy
1104
+ if is_ngram_gpu:
1105
+ ngram_gpu_new_reqs.append(req_state)
1106
1107
# Update the states of the running/resumed requests.
1108
is_last_rank = get_pp_group().is_last_rank
1109
req_data = scheduler_output.scheduled_cached_reqs
1110
scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
1111
1112
+ # Save scheduler-allocated spec lengths before trimming so
1113
+ # prev_num_draft_len keeps the optimistic count for rejection correction.
1114
+ original_num_spec_per_req: dict[str, int] = {}
1115
+ if (
1116
+ self.speculative_config is not None
1117
+ and self.speculative_config.use_ngram_gpu()
1118
+ ):
1119
+ for req_id, toks in scheduled_spec_tokens.items():
1120
+ original_num_spec_per_req[req_id] = len(toks)
1121
+ update_scheduler_for_invalid_drafts(
1122
+ self._num_valid_draft_tokens_event,
1123
+ self._num_valid_draft_tokens_cpu,
1124
+ scheduler_output,
1125
+ self.input_batch.req_id_to_index,
1126
+ )
1127
+
1128
# Wait until valid_sampled_tokens_count is copied to cpu,
1129
# then use it to update actual num_computed_tokens of each request.
1130
valid_sampled_token_count = self._get_valid_sampled_token_count()
@@ -1076,13 +1141,13 @@ class GPUModelRunner(
1141
# prev_num_draft_len is used in async scheduling mode with
1142
# spec decode. it indicates if need to update num_computed_tokens
1143
# of the request. for example:
1079
- # fist step: num_computed_tokens = 0, spec_tokens = [],
1144
+ # first step: num_computed_tokens = 0, spec_tokens = [],
1145
# prev_num_draft_len = 0.
1146
# second step: num_computed_tokens = 100(prompt length),
1147
# spec_tokens = [a,b], prev_num_draft_len = 0.
1148
# third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d],
1149
# prev_num_draft_len = 2.
1085
- # num_computed_tokens in first step and second step does't contain
1150
+ # num_computed_tokens in first step and second step doesn't contain
1151
# the spec tokens length, but in third step it contains the
1152
# spec tokens length. we only need to update num_computed_tokens
1153
# when prev_num_draft_len > 0.
@@ -1096,6 +1161,9 @@ class GPUModelRunner(
1161
num_computed_tokens -= num_rejected
1162
req_state.output_token_ids.extend([-1] * num_accepted)
1163
1164
+ if is_ngram_gpu and num_accepted > 0 and req_index is not None:
1165
+ self.input_batch.num_tokens_no_spec[req_index] += num_accepted
1166
+
1167
# Update the cached states.
1168
req_state.num_computed_tokens = num_computed_tokens
1169
@@ -1156,6 +1224,9 @@ class GPUModelRunner(
1224
req_state.output_token_ids = resumed_token_ids[-num_output_tokens:]
1225
1226
reqs_to_add.append(req_state)
1227
+ # Track resumed requests for ngram_gpu full tensor copy
1228
+ if is_ngram_gpu:
1229
+ ngram_gpu_new_reqs.append(req_state)
1230
continue
1231
1232
# Update the persistent batch.
@@ -1176,6 +1247,11 @@ class GPUModelRunner(
1247
1248
# Add spec_token_ids to token_ids_cpu.
1249
self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens)
1250
+ # Restore scheduler-side draft count after ngram trimming.
1251
+ if original_num_spec_per_req:
1252
+ orig = original_num_spec_per_req.get(req_id, 0)
1253
+ if orig != req_state.prev_num_draft_len:
1254
+ req_state.prev_num_draft_len = orig
1255
1256
# Add the new or resumed requests to the persistent batch.
1257
# The smaller empty indices are filled first.
@@ -1190,6 +1266,18 @@ class GPUModelRunner(
1266
# Refresh batch metadata with any pending updates.
1267
self.input_batch.refresh_metadata()
1268
1269
+ # Incrementally update ngram_gpu tensors after batch is stable
1270
+ if is_ngram_gpu:
1271
+ update_ngram_gpu_tensors_incremental(
1272
+ self.input_batch,
1273
+ self.token_ids_gpu_tensor,
1274
+ self.num_tokens_no_spec_gpu,
1275
+ ngram_gpu_new_reqs,
1276
+ self.device,
1277
+ _pinned_idx_buf=self._ngram_pinned_idx_buf,
1278
+ _pinned_val_buf=self._ngram_pinned_val_buf,
1279
+ )
1280
+
1281
def _update_states_after_model_execute(
1282
self, output_token_ids: torch.Tensor, scheduler_output: "SchedulerOutput"
1283
) -> None:
@@ -3412,6 +3500,23 @@ class GPUModelRunner(
3500
else:
3501
logger.error("RoutedExpertsCapturer not initialized.")
3502
3503
+ # If ngram_gpu is used, we need to copy the scheduler_output to avoid
3504
+ # the modification has influence on the scheduler_output in engine core process.
3505
+ # The replace is much faster than deepcopy.
3506
+ if (
3507
+ self.speculative_config is not None
3508
+ and self.speculative_config.use_ngram_gpu()
3509
+ ):
3510
+ num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy()
3511
+ spec_decode_tokens_copy = (
3512
+ scheduler_output.scheduled_spec_decode_tokens.copy()
3513
+ )
3514
+ scheduler_output = replace(
3515
+ scheduler_output,
3516
+ num_scheduled_tokens=num_scheduled_tokens_copy,
3517
+ scheduled_spec_decode_tokens=spec_decode_tokens_copy,
3518
+ )
3519
+
3520
if scheduler_output.preempted_req_ids and has_kv_transfer_group():
3521
get_kv_transfer_group().handle_preemptions(
3522
scheduler_output.preempted_req_ids
@@ -3825,6 +3930,32 @@ class GPUModelRunner(
3930
self._copy_valid_sampled_token_count(
3931
next_token_ids, valid_sampled_tokens_count
3932
)
3933
+ self._draft_token_ids = torch.zeros(
3934
+ 1, device=self.device, dtype=torch.int32
3935
+ ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
3936
+ self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
3937
+ elif (
3938
+ spec_config.use_ngram_gpu()
3939
+ and not spec_config.disable_padded_drafter_batch
3940
+ ):
3941
+ assert isinstance(self.drafter, NgramProposerGPU)
3942
+ sampled_token_ids = sampler_output.sampled_token_ids
3943
+ if input_fits_in_drafter:
3944
+ propose_draft_token_ids(sampled_token_ids)
3945
+ elif self.valid_sampled_token_count_event is not None:
3946
+ assert spec_decode_common_attn_metadata is not None
3947
+ next_token_ids, valid_sampled_tokens_count, _ = (
3948
+ self.drafter.update_token_ids_ngram(
3949
+ sampled_token_ids,
3950
+ self.input_batch,
3951
+ self.token_ids_gpu_tensor,
3952
+ self.num_tokens_no_spec_gpu,
3953
+ self.discard_request_mask.gpu,
3954
+ )
3955
+ )
3956
+ self._copy_valid_sampled_token_count(
3957
+ next_token_ids, valid_sampled_tokens_count
3958
+ )
3959
# Since we couldn't run the drafter,
3960
# just use zeros for the draft tokens.
3961
self._draft_token_ids = torch.zeros(
@@ -4064,6 +4195,52 @@ class GPUModelRunner(
4195
self.input_batch.token_ids_cpu,
4196
slot_mappings=slot_mappings,
4197
)
4198
+ if isinstance(self.drafter, NgramProposer):
4199
+ assert isinstance(sampled_token_ids, list), (
4200
+ "sampled_token_ids should be a python list when ngram is used."
4201
+ )
4202
+ draft_token_ids = self.drafter.propose(
4203
+ sampled_token_ids,
4204
+ self.input_batch.num_tokens_no_spec,
4205
+ self.input_batch.token_ids_cpu,
4206
+ )
4207
+ elif spec_config.use_ngram_gpu():
4208
+ assert isinstance(self.drafter, NgramProposerGPU)
4209
+ (
4210
+ next_token_ids,
4211
+ valid_sampled_tokens_count,
4212
+ valid_sampled_token_ids_gpu,
4213
+ ) = self.drafter.update_token_ids_ngram(
4214
+ sampled_token_ids,
4215
+ self.input_batch,
4216
+ self.token_ids_gpu_tensor,
4217
+ self.num_tokens_no_spec_gpu,
4218
+ self.discard_request_mask.gpu,
4219
+ )
4220
+ self._copy_valid_sampled_token_count(
4221
+ next_token_ids, valid_sampled_tokens_count
4222
+ )
4223
+
4224
+ batch_size = next_token_ids.shape[0]
4225
+
4226
+ draft_token_ids, num_valid_draft_tokens = self.drafter.propose(
4227
+ self.num_tokens_no_spec_gpu[:batch_size],
4228
+ self.token_ids_gpu_tensor[:batch_size],
4229
+ valid_sampled_token_ids_gpu,
4230
+ valid_sampled_tokens_count,
4231
+ )
4232
+
4233
+ # Cache valid draft counts for scheduler-side trimming.
4234
+ self._num_valid_draft_tokens = num_valid_draft_tokens
4235
+
4236
+ # Async D2H copy on a dedicated stream.
4237
+ copy_num_valid_draft_tokens(
4238
+ self._num_valid_draft_tokens_cpu,
4239
+ self._num_valid_draft_tokens_copy_stream,
4240
+ self._num_valid_draft_tokens_event,
4241
+ self._num_valid_draft_tokens,
4242
+ self.input_batch.num_reqs,
4243
+ )
4244
elif spec_config.method == "suffix":
4245
assert isinstance(sampled_token_ids, list)
4246
assert isinstance(self.drafter, SuffixDecodingProposer)