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

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 12, 12:15 UTC

[Model Runner V2] Implement multi-step Eagle with CUDA graph (#29559)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Nov 27, 2025 at 00:09 UTC da3222f371b48c8e2548ec22767523394580a1c5
4 files changed +526 -70
vllm/v1/worker/gpu/cudagraph_utils.py
+5 -4
@@ -233,10 +233,11 @@ def prepare_inputs_to_capture(
233 query_start_loc.np[num_reqs:] = num_tokens
234 query_start_loc.copy_to_gpu()
235 seq_lens_np = np.full(num_reqs, max_model_len, dtype=np.int32)
236 - # HACK(woosuk): To optimize warmup time, we use 1 (instead of max_model_len)
237 - # for seq_lens. This leads to a mismatch between seq_lens (GPU) and
238 - # seq_lens_np (CPU), which might cause issues in some attention backends.
239 - input_buffers.seq_lens[:num_reqs] = 1
236 + # HACK(woosuk): For faster warmup, we set seq_lens (GPU) to num_tokens
237 + # rather than max_model_len. This introduces a discrepancy between
238 + # seq_lens (on GPU) and seq_lens_np (on CPU), which may cause issues for
239 + # certain attention backends.
240 + input_buffers.seq_lens[:num_reqs] = num_tokens
241 input_buffers.seq_lens[num_reqs:] = 0
242
243 input_block_tables = [x[:num_reqs] for x in block_tables.input_block_tables]
vllm/v1/worker/gpu/model_runner.py
+19 -34
@@ -140,10 +140,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
140 self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
141
142 # CUDA graphs.
143 - self.cudagraph_manager = CudaGraphManager(
144 - vllm_config=self.vllm_config,
145 - device=self.device,
146 - )
143 + self.cudagraph_manager = CudaGraphManager(self.vllm_config, self.device)
144
145 def get_supported_tasks(self) -> tuple[str]:
146 return ("generate",)
@@ -203,6 +200,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
200 self.vllm_config,
201 self.device,
202 )
203 + if self.do_spec_decode:
204 + # HACK(woosuk)
205 + self.speculator.set_attn(
206 + self.kv_cache_config,
207 + self.attn_metadata_builders,
208 + self.block_tables,
209 + )
210 +
211 # TODO(woosuk): Support other backends.
212 if not all(b.get_name() == "FLASH_ATTN" for b in self.attn_backends.values()):
213 raise NotImplementedError("Only FLASH_ATTN backend is supported currently.")
@@ -297,35 +302,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
302 logits = self.model.compute_logits(hidden_states)
303 self.sampler(logits, sampling_metadata)
304
300 - @torch.inference_mode()
301 - def _dummy_speculator_run(
302 - self,
303 - hidden_states: torch.Tensor,
304 - aux_hidden_states: list[torch.Tensor] | None,
305 - ) -> None:
306 - num_tokens = hidden_states.shape[0]
307 - num_reqs = min(num_tokens, self.max_num_reqs)
308 - input_batch = InputBatch.make_dummy(
309 - num_reqs=num_reqs,
310 - num_tokens=num_tokens,
311 - input_buffers=self.input_buffers,
312 - device=self.device,
313 - )
314 - sampling_metadata = SamplingMetadata.make_dummy(
315 - num_reqs=num_reqs,
316 - device=self.device,
317 - )
318 - num_sampled = torch.ones(num_reqs, dtype=torch.int32, device=self.device)
319 - num_rejected = torch.zeros(num_reqs, dtype=torch.int32, device=self.device)
320 - self.propose_draft(
321 - input_batch=input_batch,
322 - sampling_metadata=sampling_metadata,
323 - last_hidden_states=hidden_states,
324 - aux_hidden_states=aux_hidden_states,
325 - num_sampled=num_sampled,
326 - num_rejected=num_rejected,
327 - )
328 -
305 @torch.inference_mode()
306 def profile_run(self) -> None:
307 hidden_states, sample_hidden_states = self._dummy_run(
@@ -334,7 +310,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
310 )
311 self._dummy_sampler_run(sample_hidden_states)
312 if self.do_spec_decode:
337 - self._dummy_speculator_run(hidden_states, None)
313 + num_tokens_across_dp = make_num_tokens_across_dp(
314 + self.dp_size, self.max_num_tokens
315 + )
316 + self.speculator.run_model(
317 + self.max_num_tokens,
318 + attn_metadata=None,
319 + num_tokens_across_dp=num_tokens_across_dp,
320 + )
321 torch.cuda.synchronize()
322 del hidden_states, sample_hidden_states
323 gc.collect()
@@ -368,6 +351,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
351 attn_metadata_builders=self.attn_metadata_builders,
352 kv_cache_config=self.kv_cache_config,
353 )
354 + if self.do_spec_decode:
355 + self.speculator.capture_model()
356
357 end_time = time.perf_counter()
358 end_free_gpu_memory = torch.cuda.mem_get_info()[0]
vllm/v1/worker/gpu/spec_decode/eagle.py
+390 -32
@@ -1,17 +1,29 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 +import numpy as np
6 import torch
7 import torch.nn as nn
8
9 from vllm.config import VllmConfig
10 from vllm.config.compilation import CUDAGraphMode
11 from vllm.forward_context import set_forward_context
12 +from vllm.logger import init_logger
13 from vllm.model_executor.model_loader import get_model
14 from vllm.triton_utils import tl, triton
11 -from vllm.v1.worker.gpu.input_batch import InputBatch
15 +from vllm.utils.platform_utils import is_pin_memory_available
16 +from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
17 +from vllm.v1.kv_cache_interface import KVCacheConfig
18 +from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
19 +from vllm.v1.worker.gpu.block_table import BlockTables
20 +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
21 from vllm.v1.worker.gpu.sampler import gumbel_sample
22 +from vllm.v1.worker.gpu.spec_decode.eagle_cudagraph import EagleCudaGraphManager
23 from vllm.v1.worker.gpu.states import SamplingMetadata
24
25 +logger = init_logger(__name__)
26 +
27
28 class EagleSpeculator:
29 def __init__(self, vllm_config: VllmConfig, device: torch.device):
@@ -27,13 +39,48 @@ class EagleSpeculator:
39 self.scheduler_config = vllm_config.scheduler_config
40 self.max_num_reqs = self.scheduler_config.max_num_seqs
41 self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
42 + self.max_model_len = vllm_config.model_config.max_model_len
43 + # We need to get the hidden size from the draft model config because
44 + # the draft model's hidden size can be different from the target model's
45 + # hidden size (e.g., Llama 3.3 70B).
46 + self.hidden_size = self.draft_model_config.get_hidden_size()
47 + self.vocab_size = self.draft_model_config.get_vocab_size()
48 + self.pin_memory = is_pin_memory_available()
49 + self.dtype = vllm_config.model_config.dtype
50
31 - self.input_ids = torch.zeros(
32 - self.max_num_tokens, dtype=torch.int32, device=device
51 + self.input_buffers = InputBuffers(
52 + max_num_reqs=self.max_num_reqs,
53 + max_num_tokens=self.max_num_tokens,
54 + hidden_size=self.hidden_size,
55 + vocab_size=self.vocab_size,
56 + dtype=self.dtype,
57 + device=device,
58 + pin_memory=self.pin_memory,
59 + )
60 + self.hidden_states = torch.zeros(
61 + self.max_num_tokens,
62 + self.hidden_size,
63 + dtype=self.dtype,
64 + device=device,
65 + )
66 + self.temperature = torch.zeros(
67 + self.max_num_reqs,
68 + dtype=torch.float32,
69 + device=device,
70 )
34 - self.positions = torch.zeros(
35 - self.max_num_tokens, dtype=torch.int64, device=device
71 + self.seeds = torch.zeros(
72 + self.max_num_reqs,
73 + dtype=torch.int64,
74 + device=device,
75 )
76 + self.draft_tokens = torch.zeros(
77 + self.max_num_reqs,
78 + self.num_speculative_steps,
79 + dtype=torch.int64,
80 + device=device,
81 + )
82 +
83 + self.cudagraph_manager = EagleCudaGraphManager(vllm_config, device)
84
85 def load_model(self, target_model: nn.Module) -> None:
86 from vllm.compilation.backends import set_model_tag
@@ -49,6 +96,91 @@ class EagleSpeculator:
96 del self.model.lm_head
97 self.model.lm_head = target_model.lm_head
98
99 + def set_attn(
100 + self,
101 + kv_cache_config: KVCacheConfig,
102 + attn_metadata_builders: list[AttentionMetadataBuilder],
103 + block_tables: BlockTables,
104 + ) -> None:
105 + self.kv_cache_config = kv_cache_config
106 + self.attn_metadata_builders = attn_metadata_builders
107 + self.block_tables = block_tables
108 +
109 + @torch.inference_mode()
110 + def run_model(
111 + self,
112 + num_tokens: int,
113 + attn_metadata: dict[str, Any],
114 + num_tokens_across_dp: torch.Tensor | None,
115 + ) -> tuple[torch.Tensor, torch.Tensor]:
116 + with set_forward_context(
117 + attn_metadata,
118 + self.vllm_config,
119 + num_tokens=num_tokens,
120 + cudagraph_runtime_mode=CUDAGraphMode.NONE,
121 + num_tokens_across_dp=num_tokens_across_dp,
122 + ):
123 + ret_hidden_states = self.model(
124 + input_ids=self.input_buffers.input_ids.gpu[:num_tokens],
125 + positions=self.input_buffers.positions[:num_tokens],
126 + hidden_states=self.hidden_states[:num_tokens],
127 + )
128 + if self.method == "mtp":
129 + last_hidden_states = ret_hidden_states
130 + hidden_states = ret_hidden_states
131 + else:
132 + last_hidden_states, hidden_states = ret_hidden_states
133 + return last_hidden_states, hidden_states
134 +
135 + def generate_draft(
136 + self,
137 + num_reqs: int,
138 + attn_metadata: dict[str, Any],
139 + num_tokens_across_dp: torch.Tensor | None,
140 + ) -> None:
141 + pos = self.input_buffers.positions[:num_reqs]
142 + query_start_loc = self.input_buffers.query_start_loc.gpu[: num_reqs + 1]
143 + for step in range(1, self.num_speculative_steps):
144 + # Run the eagle model.
145 + last_hidden_states, hidden_states = self.run_model(
146 + num_reqs, attn_metadata, num_tokens_across_dp
147 + )
148 + logits = self.model.compute_logits(last_hidden_states)
149 +
150 + # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
151 + # used for draft and target sampling.
152 + draft_tokens = gumbel_sample(
153 + logits,
154 + self.temperature[:num_reqs],
155 + self.seeds[:num_reqs],
156 + pos + 1,
157 + apply_temperature=True,
158 + )
159 + self.draft_tokens[:num_reqs, step] = draft_tokens
160 +
161 + if step < self.num_speculative_steps - 1:
162 + # Update the inputs for the next step.
163 + update_eagle_inputs(
164 + draft_tokens,
165 + hidden_states,
166 + self.input_buffers,
167 + self.hidden_states,
168 + self.max_model_len,
169 + )
170 + self.block_tables.compute_slot_mappings(query_start_loc, pos)
171 +
172 + def capture_model(self) -> None:
173 + if self.num_speculative_steps == 1:
174 + return
175 + logger.info("Capturing model for Eagle speculator...")
176 + self.cudagraph_manager.capture(
177 + self.generate_draft,
178 + self.input_buffers,
179 + self.block_tables,
180 + self.attn_metadata_builders,
181 + self.kv_cache_config,
182 + )
183 +
184 @torch.inference_mode()
185 def propose(
186 self,
@@ -80,64 +212,110 @@ class EagleSpeculator:
212 )
213 else:
214 hidden_states = last_hidden_states
215 + num_tokens = input_batch.num_tokens_after_padding
216 + self.hidden_states[:num_tokens] = hidden_states
217
218 # Get the input ids and last token indices for the speculator.
219 last_token_indices = prepare_eagle_inputs(
86 - self.input_ids,
220 + self.input_buffers,
221 input_batch,
222 num_sampled,
223 num_rejected,
224 last_sampled,
225 next_prefill_tokens,
226 )
93 - input_ids = self.input_ids[: input_batch.num_tokens_after_padding]
227
228 # Prefill: Run the eagle speculator with eager mode.
96 - with set_forward_context(
229 + # TODO(woosuk): Support CUDA graph for prefill.
230 + last_hidden_states, hidden_states = self.run_model(
231 + num_tokens,
232 input_batch.attn_metadata,
98 - self.vllm_config,
99 - num_tokens=input_batch.num_tokens_after_padding,
100 - cudagraph_runtime_mode=CUDAGraphMode.NONE,
101 - ):
102 - ret_hidden_states = self.model(
103 - input_ids=input_ids,
104 - positions=input_batch.positions,
105 - hidden_states=hidden_states,
106 - )
107 - if self.method == "mtp":
108 - last_hidden_states = ret_hidden_states
109 - hidden_states = ret_hidden_states
110 - else:
111 - last_hidden_states, hidden_states = ret_hidden_states
233 + num_tokens_across_dp=None, # FIXME
234 + )
235 sample_hidden_states = last_hidden_states[last_token_indices]
236 logits = self.model.compute_logits(sample_hidden_states)
237
238 num_reqs = input_batch.num_reqs
239 cu_num_logits = input_batch.cu_num_logits[:num_reqs]
117 - temperature = sampling_metadata.temperature[cu_num_logits]
118 - seed = sampling_metadata.seeds[cu_num_logits]
119 - # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
120 - # used for draft and target sampling.
121 - pos = input_batch.positions[last_token_indices] + 1
240 # NOTE(woosuk): For draft sampling, we only consider the temperature
241 # and ignore the other sampling parameters such as top_k and top_p,
242 # for simplicity and performance.
243 # While this may slightly degrade the acceptance rate, it does not
244 # affect the output distribution after rejection sampling.
245 + temperature = self.temperature[:num_reqs]
246 + seeds = self.seeds[:num_reqs]
247 + pos = self.input_buffers.positions[:num_reqs]
248 + # Gather the values and copy them to the pre-allocated buffers.
249 + torch.gather(sampling_metadata.temperature, 0, cu_num_logits, out=temperature)
250 + torch.gather(sampling_metadata.seeds, 0, cu_num_logits, out=seeds)
251 + torch.gather(input_batch.positions, 0, last_token_indices, out=pos)
252 + # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
253 + # used for draft and target sampling.
254 draft_tokens = gumbel_sample(
128 - logits, temperature, seed, pos, apply_temperature=True
255 + logits, temperature, seeds, pos + 1, apply_temperature=True
256 )
257 if self.num_speculative_steps == 1:
258 # Early exit.
259 return draft_tokens.view(-1, 1)
133 - raise NotImplementedError("num_speculative_steps > 1 is not supported yet.")
260 +
261 + # Save the draft tokens for the first step.
262 + self.draft_tokens[:num_reqs, 0] = draft_tokens
263 + # Prepare the inputs for the decode steps.
264 + prepare_eagle_decode(
265 + draft_tokens,
266 + hidden_states,
267 + last_token_indices,
268 + input_batch.seq_lens,
269 + num_rejected,
270 + self.input_buffers,
271 + self.hidden_states,
272 + self.max_model_len,
273 + self.max_num_reqs,
274 + )
275 + query_start_loc = self.input_buffers.query_start_loc
276 + query_start_loc_gpu = query_start_loc.gpu[: num_reqs + 1]
277 + slot_mappings = self.block_tables.compute_slot_mappings(
278 + query_start_loc_gpu, pos
279 + )
280 +
281 + cudagraph_size = self.cudagraph_manager.get_cudagraph_size(num_reqs)
282 + if cudagraph_size is not None:
283 + # Run CUDA graph.
284 + self.cudagraph_manager.run(cudagraph_size)
285 + return self.draft_tokens[:num_reqs]
286 +
287 + # Run eager mode.
288 + query_start_loc.np[: num_reqs + 1] = np.arange(num_reqs + 1)
289 + query_start_loc_cpu = query_start_loc.cpu[: num_reqs + 1]
290 + # HACK(woosuk)
291 + seq_lens_np = np.full(num_reqs, self.max_model_len, dtype=np.int32)
292 + block_tables = [x[:num_reqs] for x in self.block_tables.input_block_tables]
293 +
294 + # FIXME(woosuk): This is UNSAFE!!
295 + attn_metadata = build_attn_metadata(
296 + attn_metadata_builders=self.attn_metadata_builders,
297 + num_reqs=num_reqs,
298 + num_tokens=num_reqs,
299 + query_start_loc_gpu=query_start_loc_gpu,
300 + query_start_loc_cpu=query_start_loc_cpu,
301 + seq_lens=self.input_buffers.seq_lens[:num_reqs],
302 + seq_lens_np=seq_lens_np,
303 + num_computed_tokens_cpu=None, # FIXME
304 + block_tables=block_tables,
305 + slot_mappings=slot_mappings,
306 + kv_cache_config=self.kv_cache_config,
307 + )
308 + self.generate_draft(num_reqs, attn_metadata, num_tokens_across_dp=None) # FIXME
309 + return self.draft_tokens[:num_reqs]
310
311
312 @triton.jit
313 def _prepare_eagle_inputs_kernel(
314 last_token_indices_ptr,
315 eagle_input_ids_ptr,
316 + eagle_positions_ptr,
317 target_input_ids_ptr,
318 + target_positions_ptr,
319 idx_mapping_ptr,
320 last_sampled_ptr,
321 next_prefill_tokens_ptr,
@@ -175,9 +353,16 @@ def _prepare_eagle_inputs_kernel(
353 tl.store(last_token_indices_ptr + batch_idx, last_token_index)
354 tl.store(eagle_input_ids_ptr + last_token_index, next_token)
355
356 + # Copy positions.
357 + for i in range(0, query_len, BLOCK_SIZE):
358 + block = i + tl.arange(0, BLOCK_SIZE)
359 + mask = block < query_len
360 + target_pos = tl.load(target_positions_ptr + query_start + block, mask=mask)
361 + tl.store(eagle_positions_ptr + query_start + block, target_pos, mask=mask)
362 +
363
364 def prepare_eagle_inputs(
180 - eagle_input_ids: torch.Tensor,
365 + input_buffers: InputBuffers,
366 input_batch: InputBatch,
367 # [num_reqs]
368 num_sampled: torch.Tensor,
@@ -192,12 +377,14 @@ def prepare_eagle_inputs(
377 last_token_indices = torch.empty(
378 num_reqs,
379 dtype=torch.int64,
195 - device=eagle_input_ids.device,
380 + device=num_sampled.device,
381 )
382 _prepare_eagle_inputs_kernel[(num_reqs,)](
383 last_token_indices,
199 - eagle_input_ids,
384 + input_buffers.input_ids.gpu,
385 + input_buffers.positions,
386 input_batch.input_ids,
387 + input_batch.positions,
388 input_batch.idx_mapping,
389 last_sampled,
390 next_prefill_tokens,
@@ -207,3 +394,174 @@ def prepare_eagle_inputs(
394 BLOCK_SIZE=1024,
395 )
396 return last_token_indices
397 +
398 +
399 +@triton.jit
400 +def _prepare_eagle_docode_kernel(
401 + draft_tokens_ptr,
402 + output_hidden_states_ptr,
403 + output_hidden_states_stride,
404 + last_token_indices_ptr,
405 + target_seq_lens_ptr,
406 + num_rejected_ptr,
407 + input_ids_ptr,
408 + positions_ptr,
409 + input_hidden_states_ptr,
410 + input_hidden_states_stride,
411 + query_start_loc_ptr,
412 + seq_lens_ptr,
413 + hidden_size,
414 + max_model_len,
415 + max_num_reqs,
416 + BLOCK_SIZE: tl.constexpr,
417 +):
418 + req_idx = tl.program_id(0)
419 + num_reqs = tl.num_programs(0) - 1
420 + if req_idx == num_reqs:
421 + # Compute query_start_loc. Pad it with the last query_start_loc
422 + # for CUDA graphs.
423 + for i in range(0, max_num_reqs + 1, BLOCK_SIZE):
424 + block = i + tl.arange(0, BLOCK_SIZE)
425 + q = tl.where(block < num_reqs, block, num_reqs)
426 + mask = block < max_num_reqs + 1
427 + tl.store(query_start_loc_ptr + block, q, mask=mask)
428 + # Pad seq_lens for CUDA graphs.
429 + for i in range(req_idx, max_num_reqs, BLOCK_SIZE):
430 + block = i + tl.arange(0, BLOCK_SIZE)
431 + mask = block < max_num_reqs
432 + tl.store(seq_lens_ptr + block, 0, mask=mask)
433 + return
434 +
435 + # draft token -> input id.
436 + draft_token = tl.load(draft_tokens_ptr + req_idx)
437 + tl.store(input_ids_ptr + req_idx, draft_token)
438 +
439 + # output hidden states -> input hidden states.
440 + src_idx = tl.load(last_token_indices_ptr + req_idx)
441 + for i in range(0, hidden_size, BLOCK_SIZE):
442 + block = i + tl.arange(0, BLOCK_SIZE)
443 + mask = block < hidden_size
444 + output_hidden_states = tl.load(
445 + output_hidden_states_ptr + src_idx * output_hidden_states_stride + block,
446 + mask=mask,
447 + )
448 + tl.store(
449 + input_hidden_states_ptr + req_idx * input_hidden_states_stride + block,
450 + output_hidden_states,
451 + mask=mask,
452 + )
453 +
454 + # Compute position and seq_lens.
455 + # NOTE(woosuk): To prevent out-of-range access, we clamp these values
456 + # if they reach the max model length.
457 + position = tl.load(positions_ptr + req_idx)
458 + position = tl.minimum(position + 1, max_model_len - 1)
459 + tl.store(positions_ptr + req_idx, position)
460 +
461 + target_seq_len = tl.load(target_seq_lens_ptr + req_idx)
462 + num_rejected = tl.load(num_rejected_ptr + req_idx)
463 + seq_len = target_seq_len - num_rejected
464 + seq_len = tl.minimum(seq_len + 1, max_model_len)
465 + tl.store(seq_lens_ptr + req_idx, seq_len)
466 +
467 +
468 +def prepare_eagle_decode(
469 + draft_tokens: torch.Tensor,
470 + output_hidden_states: torch.Tensor,
471 + last_token_indices: torch.Tensor,
472 + target_seq_lens: torch.Tensor,
473 + num_rejected: torch.Tensor,
474 + input_buffers: InputBuffers,
475 + input_hidden_states: torch.Tensor,
476 + max_model_len: int,
477 + max_num_reqs: int,
478 +):
479 + num_reqs = draft_tokens.shape[0]
480 + hidden_size = output_hidden_states.shape[-1]
481 + _prepare_eagle_docode_kernel[(num_reqs + 1,)](
482 + draft_tokens,
483 + output_hidden_states,
484 + output_hidden_states.stride(0),
485 + last_token_indices,
486 + target_seq_lens,
487 + num_rejected,
488 + input_buffers.input_ids.gpu,
489 + input_buffers.positions,
490 + input_hidden_states,
491 + input_hidden_states.stride(0),
492 + input_buffers.query_start_loc.gpu,
493 + input_buffers.seq_lens,
494 + hidden_size,
495 + max_model_len,
496 + max_num_reqs,
497 + BLOCK_SIZE=1024,
498 + )
499 +
500 +
501 +@triton.jit
502 +def _update_eagle_inputs_kernel(
503 + input_ids_ptr,
504 + positions_ptr,
505 + input_hidden_states_ptr,
506 + input_hidden_states_stride,
507 + seq_lens_ptr,
508 + max_model_len,
509 + draft_tokens_ptr,
510 + output_hidden_states_ptr,
511 + output_hidden_states_stride,
512 + hidden_size,
513 + BLOCK_SIZE: tl.constexpr,
514 +):
515 + req_idx = tl.program_id(0)
516 +
517 + # Draft token -> Input ID.
518 + draft_token = tl.load(draft_tokens_ptr + req_idx)
519 + tl.store(input_ids_ptr + req_idx, draft_token)
520 +
521 + # Output hidden states -> Input hidden states.
522 + for i in range(0, hidden_size, BLOCK_SIZE):
523 + block = i + tl.arange(0, BLOCK_SIZE)
524 + mask = block < hidden_size
525 + output_hidden_states = tl.load(
526 + output_hidden_states_ptr + req_idx * output_hidden_states_stride + block,
527 + mask=mask,
528 + )
529 + tl.store(
530 + input_hidden_states_ptr + req_idx * input_hidden_states_stride + block,
531 + output_hidden_states,
532 + mask=mask,
533 + )
534 +
535 + # Increment position and seq_lens.
536 + # NOTE(woosuk): To prevent out-of-range access, we clamp these values
537 + # if they reach the max model length.
538 + position = tl.load(positions_ptr + req_idx)
539 + position = tl.minimum(position + 1, max_model_len - 1)
540 + tl.store(positions_ptr + req_idx, position)
541 +
542 + seq_len = tl.load(seq_lens_ptr + req_idx)
543 + seq_len = tl.minimum(seq_len + 1, max_model_len)
544 + tl.store(seq_lens_ptr + req_idx, seq_len)
545 +
546 +
547 +def update_eagle_inputs(
548 + draft_tokens: torch.Tensor,
549 + output_hidden_states: torch.Tensor,
550 + input_buffers: InputBuffers,
551 + hidden_states: torch.Tensor,
552 + max_model_len: int,
553 +):
554 + num_reqs, hidden_size = output_hidden_states.shape
555 + _update_eagle_inputs_kernel[(num_reqs,)](
556 + input_buffers.input_ids.gpu,
557 + input_buffers.positions,
558 + hidden_states,
559 + hidden_states.stride(0),
560 + input_buffers.seq_lens,
561 + max_model_len,
562 + draft_tokens,
563 + output_hidden_states,
564 + output_hidden_states.stride(0),
565 + hidden_size,
566 + BLOCK_SIZE=1024,
567 + )
vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py new
+112
@@ -0,0 +1,112 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable
4 +
5 +import torch
6 +
7 +from vllm.config import VllmConfig
8 +from vllm.config.compilation import CUDAGraphMode
9 +from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
10 +from vllm.v1.kv_cache_interface import KVCacheConfig
11 +from vllm.v1.worker.gpu.block_table import BlockTables
12 +from vllm.v1.worker.gpu.cudagraph_utils import (
13 + capture_graphs,
14 + get_cudagraph_sizes,
15 + prepare_inputs_to_capture,
16 +)
17 +from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
18 +from vllm.v1.worker.gpu.input_batch import InputBuffers
19 +
20 +
21 +class EagleCudaGraphManager:
22 + def __init__(
23 + self,
24 + vllm_config: VllmConfig,
25 + device: torch.device,
26 + ):
27 + self.vllm_config = vllm_config
28 + self.scheduler_config = vllm_config.scheduler_config
29 + self.device = device
30 +
31 + self.max_model_len = vllm_config.model_config.max_model_len
32 + self.max_num_reqs = self.scheduler_config.max_num_seqs
33 + self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
34 + self.dp_size = vllm_config.parallel_config.data_parallel_size
35 + self.compilation_config = vllm_config.compilation_config
36 + assert self.compilation_config is not None
37 +
38 + if self.compilation_config.cudagraph_mode is None:
39 + self.cudagraph_mode = CUDAGraphMode.NONE
40 + else:
41 + self.cudagraph_mode = self.compilation_config.cudagraph_mode
42 + if self.cudagraph_mode == CUDAGraphMode.FULL:
43 + # NOTE(woosuk): For Eagle, we only use CUDA graphs for decode.
44 + self.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
45 +
46 + self.cudagraph_sizes = get_cudagraph_sizes(
47 + self.compilation_config.cudagraph_capture_sizes,
48 + self.max_num_reqs,
49 + self.max_num_tokens,
50 + self.cudagraph_mode,
51 + )
52 +
53 + self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
54 + self.pool = torch.cuda.graph_pool_handle()
55 +
56 + def get_cudagraph_size(self, num_tokens: int) -> int | None:
57 + return self.cudagraph_sizes.get(num_tokens)
58 +
59 + def capture_graph(
60 + self,
61 + num_tokens: int,
62 + generate_fn: Callable,
63 + input_buffers: InputBuffers,
64 + block_tables: BlockTables,
65 + attn_metadata_builders: list[AttentionMetadataBuilder],
66 + kv_cache_config: KVCacheConfig,
67 + ) -> None:
68 + num_reqs = min(num_tokens, self.max_num_reqs)
69 + attn_metadata = prepare_inputs_to_capture(
70 + num_reqs,
71 + num_tokens,
72 + input_buffers,
73 + block_tables,
74 + attn_metadata_builders,
75 + self.max_model_len,
76 + kv_cache_config,
77 + )
78 + num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)
79 +
80 + # Warm up.
81 + generate_fn(num_tokens, attn_metadata, num_tokens_across_dp)
82 +
83 + # Capture the graph.
84 + assert num_tokens not in self.graphs
85 + graph = torch.cuda.CUDAGraph()
86 + with torch.cuda.graph(graph, self.pool):
87 + generate_fn(num_tokens, attn_metadata, num_tokens_across_dp)
88 + self.graphs[num_tokens] = graph
89 +
90 + @torch.inference_mode()
91 + def capture(
92 + self,
93 + generate_fn: Callable,
94 + input_buffers: InputBuffers,
95 + block_tables: BlockTables,
96 + attn_metadata_builders: list[AttentionMetadataBuilder],
97 + kv_cache_config: KVCacheConfig,
98 + ) -> None:
99 + capture_graphs(
100 + self.cudagraph_sizes,
101 + self.device,
102 + self.capture_graph,
103 + generate_fn=generate_fn,
104 + input_buffers=input_buffers,
105 + block_tables=block_tables,
106 + attn_metadata_builders=attn_metadata_builders,
107 + kv_cache_config=kv_cache_config,
108 + )
109 +
110 + def run(self, num_tokens: int) -> None:
111 + assert num_tokens in self.graphs
112 + self.graphs[num_tokens].replay()