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 5, 06:00 UTC
[Model Runner V2] Add full cuda graph support for eagle prefill (#37588)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Apr 13, 2026 at 16:01 UTC
ccf90ba784900324e8ff4e6495d13b47b27b941a
3 files changed
+210
-134
vllm/v1/worker/gpu/model_runner.py
-6
@@ -464,7 +464,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
464
slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
465
hidden_states = self.execute_model_state.hidden_states
466
aux_hidden_states = self.execute_model_state.aux_hidden_states
467
- num_tokens_across_dp = self.execute_model_state.num_tokens_across_dp
467
self.execute_model_state = None
468
469
# dummy run the eagle speculator's propose to ensure DP/EP sync.
@@ -496,7 +495,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
495
next_prefill_tokens=self.req_states.next_prefill_tokens,
496
temperature=self.sampler.sampling_states.temperature.gpu,
497
seeds=self.sampler.sampling_states.seeds.gpu,
499
- num_tokens_across_dp=num_tokens_across_dp,
498
dummy_run=True,
499
skip_attn_for_dummy_run=skip_attn,
500
mm_inputs=mm_inputs,
@@ -1110,7 +1108,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1108
hidden_states=hidden_states,
1109
aux_hidden_states=aux_hidden_states,
1110
kv_connector_output=kv_connector_output,
1113
- num_tokens_across_dp=num_tokens_across_dp,
1111
)
1112
1113
if not self.is_last_pp_rank:
@@ -1135,7 +1132,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1132
hidden_states = self.execute_model_state.hidden_states
1133
aux_hidden_states = self.execute_model_state.aux_hidden_states
1134
kv_connector_output = self.execute_model_state.kv_connector_output
1138
- num_tokens_across_dp = self.execute_model_state.num_tokens_across_dp
1135
self.execute_model_state = None
1136
1137
if not self.is_last_pp_rank:
@@ -1228,7 +1224,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1224
self.req_states.next_prefill_tokens,
1225
self.sampler.sampling_states.temperature.gpu,
1226
self.sampler.sampling_states.seeds.gpu,
1231
- num_tokens_across_dp=num_tokens_across_dp,
1227
mm_inputs=mm_inputs,
1228
)
1229
self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
@@ -1336,4 +1331,3 @@ class ExecuteModelState(NamedTuple):
1331
hidden_states: torch.Tensor | None
1332
aux_hidden_states: list[torch.Tensor] | None
1333
kv_connector_output: KVConnectorOutput | None
1339
- num_tokens_across_dp: torch.Tensor | None
vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
+6
-16
@@ -19,21 +19,16 @@ from vllm.v1.worker.utils import AttentionGroup
19
20
21
class EagleCudaGraphManager(CudaGraphManager):
22
- """CudaGraphManager for Eagle speculative decoding (FULL mode only)."""
22
+ """CudaGraphManager for Eagle speculative decoding."""
23
24
def __init__(
25
self,
26
vllm_config: VllmConfig,
27
device: torch.device,
28
cudagraph_mode: CUDAGraphMode,
29
- draft_tokens: torch.Tensor,
29
+ decode_query_len: int,
30
):
31
- assert not cudagraph_mode.has_mode(CUDAGraphMode.PIECEWISE), (
32
- "EagleCudaGraphManager does not support PIECEWISE mode yet"
33
- )
34
- # Eagle always uses uniform decode with query_len=1
35
- super().__init__(vllm_config, device, cudagraph_mode, decode_query_len=1)
36
- self.draft_tokens = draft_tokens
31
+ super().__init__(vllm_config, device, cudagraph_mode, decode_query_len)
32
33
# Use a dedicated pool for Eagle to avoid memory overlap with the main
34
# model's cudagraph. The base class uses a shared global pool, but Eagle's
@@ -44,7 +39,7 @@ class EagleCudaGraphManager(CudaGraphManager):
39
40
def capture(
41
self,
47
- generate_fn: Callable,
42
+ forward_fn: Callable,
43
model_state: ModelState,
44
input_buffers: InputBuffers,
45
block_tables: BlockTables,
@@ -52,7 +47,7 @@ class EagleCudaGraphManager(CudaGraphManager):
47
kv_cache_config: KVCacheConfig,
48
progress_bar_desc: str = "Capturing CUDA graphs",
49
) -> None:
55
- """Capture CUDA graphs for Eagle speculative decoding (FULL mode only)."""
50
+ """Capture CUDA graphs for Eagle."""
51
52
def create_forward_fn(
53
desc: BatchExecutionDescriptor,
@@ -74,7 +69,7 @@ class EagleCudaGraphManager(CudaGraphManager):
69
kv_cache_config,
70
)
71
77
- return lambda cg_mode: generate_fn(
72
+ return lambda cg_mode: forward_fn(
73
num_reqs,
74
num_tokens,
75
attn_metadata,
@@ -84,8 +79,3 @@ class EagleCudaGraphManager(CudaGraphManager):
79
)
80
81
super().capture(create_forward_fn, progress_bar_desc)
87
-
88
- def run_fullgraph(self, desc: BatchExecutionDescriptor) -> torch.Tensor:
89
- """Replay a captured FULL cudagraph and return draft tokens."""
90
- super().run_fullgraph(desc)
91
- return self.draft_tokens
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+204
-112
@@ -19,11 +19,16 @@ from vllm.v1.worker.gpu.attn_utils import (
19
init_attn_backend,
20
)
21
from vllm.v1.worker.gpu.block_table import BlockTables
22
+from vllm.v1.worker.gpu.cudagraph_utils import (
23
+ get_uniform_token_count,
24
+)
25
from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
26
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
27
from vllm.v1.worker.gpu.model_states.interface import ModelState
28
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
26
-from vllm.v1.worker.gpu.spec_decode.eagle.cudagraph import EagleCudaGraphManager
29
+from vllm.v1.worker.gpu.spec_decode.eagle.cudagraph import (
30
+ EagleCudaGraphManager,
31
+)
32
from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
33
34
logger = init_logger(__name__)
@@ -76,6 +81,9 @@ class EagleSpeculator:
81
dtype=torch.int64,
82
device=device,
83
)
84
+ self.last_token_indices = torch.zeros(
85
+ self.max_num_reqs, dtype=torch.int64, device=device
86
+ )
87
88
self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
89
self.draft_model_config
@@ -95,20 +103,30 @@ class EagleSpeculator:
103
device=device,
104
)
105
98
- self.cudagraph_manager: EagleCudaGraphManager | None = None
106
+ self.prefill_cudagraph_manager: EagleCudaGraphManager | None = None
107
+ self.decode_cudagraph_manager: EagleCudaGraphManager | None = None
108
109
def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
101
- if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL:
102
- cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
103
- else:
104
- cudagraph_mode = CUDAGraphMode.NONE
105
-
106
- self.cudagraph_manager = EagleCudaGraphManager(
110
+ cudagraph_mode = self.vllm_config.compilation_config.cudagraph_mode
111
+ # Initialize cudagraph manager for draft prefill (draft position 0).
112
+ self.prefill_cudagraph_manager = EagleCudaGraphManager(
113
self.vllm_config,
114
self.device,
115
cudagraph_mode,
110
- self.draft_tokens,
116
+ self.num_speculative_steps + 1,
117
)
118
+ # Initialize cudagraph manager for draft generation (draft positions > 0).
119
+ self.decode_cudagraph_manager = EagleCudaGraphManager(
120
+ self.vllm_config,
121
+ self.device,
122
+ # Only use FULL graph mode, if available, because draft decodes
123
+ # only consist of a single token.
124
+ cudagraph_mode.decode_mode(),
125
+ decode_query_len=1,
126
+ )
127
+ # Share a single pool between prefill and decode since they never
128
+ # execute concurrently.
129
+ self.decode_cudagraph_manager.pool = self.prefill_cudagraph_manager.pool
130
131
def load_model(self, target_model: nn.Module) -> None:
132
target_attn_layer_names = get_layers_from_vllm_config(
@@ -189,6 +207,47 @@ class EagleSpeculator:
207
last_hidden_states, hidden_states = ret_hidden_states
208
return last_hidden_states, hidden_states
209
210
+ def prefill(
211
+ self,
212
+ num_reqs: int,
213
+ num_tokens: int,
214
+ attn_metadata: dict[str, Any] | None,
215
+ slot_mappings: dict[str, torch.Tensor] | None,
216
+ num_tokens_across_dp: torch.Tensor | None,
217
+ cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
218
+ mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
219
+ ) -> None:
220
+ last_token_indices = self.last_token_indices[:num_reqs]
221
+ pos = self.input_buffers.positions[last_token_indices]
222
+ idx_mapping = self.idx_mapping[:num_reqs]
223
+
224
+ last_hidden_states, hidden_states = self.run_model(
225
+ num_tokens,
226
+ attn_metadata,
227
+ slot_mappings,
228
+ num_tokens_across_dp=num_tokens_across_dp,
229
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
230
+ mm_inputs=mm_inputs,
231
+ )
232
+ sample_hidden_states = last_hidden_states[last_token_indices]
233
+ logits = self.model.compute_logits(sample_hidden_states)
234
+
235
+ # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
236
+ # used for draft and target sampling.
237
+ self.draft_tokens[:num_reqs, 0] = gumbel_sample(
238
+ logits,
239
+ idx_mapping,
240
+ self.temperature,
241
+ self.seeds,
242
+ pos + 1,
243
+ apply_temperature=True,
244
+ processed_logits_out=self.draft_logits[:, 0]
245
+ if self.draft_logits is not None
246
+ else None,
247
+ )
248
+ self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
249
+ self.input_buffers.positions[:num_reqs] = pos
250
+
251
def generate_draft(
252
self,
253
num_reqs: int,
@@ -281,19 +340,46 @@ class EagleSpeculator:
340
return attn_metadata
341
342
def capture_model(self) -> None:
284
- assert self.cudagraph_manager is not None
343
+ logger.info("Capturing model for Eagle speculator...")
344
+ # Reset indices to zeros to prevent stale values from prior
345
+ # dummy runs to cause out-of-bounds indexing during capture.
346
+ self.last_token_indices.zero_()
347
+
348
+ # Capture the prefill routine (model forward + compute_logits +
349
+ # gumbel_sample).
350
+ # For FULL graphs, the entire routine is recorded as one graph.
351
+ # For PIECEWISE, only the model's compiled regions are captured
352
+ # and the rest (compute_logits, gumbel_sample) runs eagerly.
353
+ assert self.prefill_cudagraph_manager is not None
354
+ self.prefill_cudagraph_manager.capture(
355
+ self.prefill,
356
+ self.model_state,
357
+ self.input_buffers,
358
+ self.block_tables,
359
+ self.attn_groups,
360
+ self.kv_cache_config,
361
+ progress_bar_desc="Capturing eagle prefill CUDA graphs",
362
+ )
363
+
364
if self.num_speculative_steps == 1:
365
return
366
288
- logger.info("Capturing model for Eagle speculator...")
289
- self.cudagraph_manager.capture(
367
+ # Capture the decode draft generation loop (model forward +
368
+ # compute_logits + gumbel_sample + update_eagle_inputs, for
369
+ # each step).
370
+ # For FULL graphs, the entire multi-step loop is recorded as
371
+ # one graph. For PIECEWISE, only the model's compiled regions
372
+ # are captured, and the rest (compute_logits, gumbel_sample,
373
+ # update_eagle_inputs) runs eagerly.
374
+ assert self.decode_cudagraph_manager is not None
375
+ self.decode_cudagraph_manager.capture(
376
self.generate_draft,
377
self.model_state,
378
self.input_buffers,
379
self.block_tables,
380
self.attn_groups,
381
self.kv_cache_config,
296
- progress_bar_desc="Capturing eagle CUDA graphs",
382
+ progress_bar_desc="Capturing eagle decode CUDA graphs",
383
)
384
385
@torch.inference_mode()
@@ -324,6 +410,10 @@ class EagleSpeculator:
410
mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
411
is_profile: bool = False,
412
) -> torch.Tensor:
413
+ num_tokens = input_batch.num_tokens_after_padding
414
+ num_reqs = input_batch.num_reqs
415
+ max_query_len = input_batch.num_scheduled_tokens.max()
416
+
417
# NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
418
# number of rejected tokens, we maintain the size of eagle's input_ids and
419
# hidden_states the same as the target model's. This means, we pad each
@@ -337,82 +427,88 @@ class EagleSpeculator:
427
)
428
else:
429
hidden_states = last_hidden_states
340
- num_tokens = input_batch.num_tokens_after_padding
341
- self.hidden_states[:num_tokens] = hidden_states
430
+ self.hidden_states[:num_tokens].copy_(hidden_states)
431
+
432
+ # Copy temperature, seeds, and idx mapping to the pre-allocated buffers.
433
+ # NOTE(woosuk): For draft sampling, we only consider the temperature
434
+ # and ignore the other sampling parameters such as top_k and top_p,
435
+ # for simplicity and performance.
436
+ # While this may slightly degrade the acceptance rate, it does not
437
+ # affect the output distribution after rejection sampling.
438
+ self.temperature.copy_(temperature)
439
+ self.seeds.copy_(seeds)
440
+ self.idx_mapping[:num_reqs].copy_(input_batch.idx_mapping)
441
442
# Get the input ids and last token indices for the speculator.
344
- last_token_indices = prepare_eagle_inputs(
443
+ prepare_eagle_inputs(
444
self.input_buffers,
445
input_batch,
446
+ self.last_token_indices,
447
num_sampled,
448
num_rejected,
449
last_sampled,
450
next_prefill_tokens,
451
+ self.max_num_reqs,
452
)
453
353
- # Prefill: Run the eagle speculator with eager mode.
354
- # TODO(woosuk): Support CUDA graph for prefill.
355
- last_hidden_states, hidden_states = self.run_model(
454
+ # When all requests are decoding (no true prefills), each has
455
+ # num_speculative_steps + 1 tokens, enabling FULL graph replay.
456
+ # Mixed or prefill-only batches fall back to PIECEWISE.
457
+ prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
458
+ self.prefill_cudagraph_manager,
459
+ num_reqs,
460
num_tokens,
357
- attn_metadata,
358
- slot_mappings,
359
- num_tokens_across_dp=num_tokens_across_dp,
360
- mm_inputs=mm_inputs,
461
+ get_uniform_token_count(num_reqs, num_tokens, max_query_len),
462
+ dp_size=self.dp_size,
463
+ dp_rank=self.dp_rank,
464
+ need_eager=is_profile,
465
)
362
- sample_hidden_states = last_hidden_states[last_token_indices]
363
- logits = self.model.compute_logits(sample_hidden_states)
466
365
- num_reqs = input_batch.num_reqs
366
- # NOTE(woosuk): For draft sampling, we only consider the temperature
367
- # and ignore the other sampling parameters such as top_k and top_p,
368
- # for simplicity and performance.
369
- # While this may slightly degrade the acceptance rate, it does not
370
- # affect the output distribution after rejection sampling.
371
- idx_mapping = self.idx_mapping[:num_reqs]
372
- idx_mapping.copy_(input_batch.idx_mapping)
373
- self.temperature.copy_(temperature)
374
- self.seeds.copy_(seeds)
375
-
376
- # Gather the values and copy them to the pre-allocated buffers.
377
- pos = self.input_buffers.positions[:num_reqs]
378
- torch.gather(input_batch.positions, 0, last_token_indices, out=pos)
379
- # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
380
- # used for draft and target sampling.
381
- draft_tokens = gumbel_sample(
382
- logits,
383
- idx_mapping,
384
- self.temperature,
385
- self.seeds,
386
- pos + 1,
387
- apply_temperature=True,
388
- processed_logits_out=self.draft_logits[:, 0]
389
- if self.draft_logits is not None
390
- else None,
391
- )
467
+ if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL:
468
+ # It is necessary to rebuild the attention metadata when
469
+ # replaying the FULL graph so that any attention metadata
470
+ # builder state is updated.
471
+ self._build_draft_attn_metadata(
472
+ num_reqs=num_reqs,
473
+ num_reqs_padded=prefill_batch_desc.num_reqs or num_reqs,
474
+ num_tokens_padded=prefill_batch_desc.num_tokens,
475
+ max_query_len=self.num_speculative_steps + 1,
476
+ )
477
+ # Replay the full graph for draft prefill.
478
+ assert self.prefill_cudagraph_manager is not None
479
+ self.prefill_cudagraph_manager.run_fullgraph(prefill_batch_desc)
480
+ else:
481
+ # The target model's attention metadata and slot mappings
482
+ # can directly be used for draft prefill, because of the
483
+ # identical batch shape and KV cache layout.
484
+ self.prefill(
485
+ num_reqs,
486
+ prefill_batch_desc.num_tokens,
487
+ attn_metadata,
488
+ slot_mappings,
489
+ num_tokens_across_dp=num_tokens_across_dp,
490
+ cudagraph_runtime_mode=prefill_batch_desc.cg_mode,
491
+ mm_inputs=mm_inputs,
492
+ )
493
494
if self.num_speculative_steps == 1:
495
# Early exit.
395
- return draft_tokens.view(-1, 1)
496
+ return self.draft_tokens[:num_reqs, :1]
497
397
- # Save the draft tokens for the first step.
398
- self.draft_tokens[:num_reqs, 0] = draft_tokens
498
# Prepare the inputs for the decode steps.
499
prepare_eagle_decode(
401
- draft_tokens,
402
- hidden_states,
403
- last_token_indices,
500
+ self.draft_tokens[:num_reqs, 0],
501
input_batch.seq_lens,
502
num_rejected,
503
self.input_buffers,
407
- self.hidden_states,
504
self.max_model_len,
505
self.max_num_reqs,
506
)
507
412
- # Each request produces exactly 1 token per draft decode step,
413
- # enabling FULL cudagraph.
508
+ # Each request produces exactly 1 token per draft generation step,
509
+ # enabling FULL graph replay.
510
decode_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
415
- self.cudagraph_manager,
511
+ self.decode_cudagraph_manager,
512
num_reqs,
513
num_reqs,
514
uniform_token_count=1,
@@ -426,12 +522,12 @@ class EagleSpeculator:
522
if not (dummy_run and skip_attn_for_dummy_run):
523
# Build attention metadata and slot mappings for the draft
524
# decode steps. It is necessary to rebuild the attention
429
- # metadata even when replaying the FULL cudagraph so that
430
- # any attention metadata builder state is updated.
525
+ # metadata even when replaying the FULL graph so that any
526
+ # attention metadata builder state is updated.
527
slot_mappings = self.block_tables.compute_slot_mappings(
432
- idx_mapping,
528
+ self.idx_mapping[:num_reqs],
529
self.input_buffers.query_start_loc[: num_reqs + 1],
434
- pos,
530
+ self.input_buffers.positions[:num_reqs],
531
decode_batch_desc.num_tokens,
532
)
533
slot_mappings_updated = build_slot_mappings_by_layer(
@@ -445,8 +541,9 @@ class EagleSpeculator:
541
)
542
543
if decode_batch_desc.cg_mode == CUDAGraphMode.FULL:
448
- assert self.cudagraph_manager is not None
449
- self.cudagraph_manager.run_fullgraph(decode_batch_desc)
544
+ # Replay the full graph for draft generation.
545
+ assert self.decode_cudagraph_manager is not None
546
+ self.decode_cudagraph_manager.run_fullgraph(decode_batch_desc)
547
else:
548
self.generate_draft(
549
num_reqs,
@@ -464,6 +561,8 @@ def _prepare_eagle_inputs_kernel(
561
last_token_indices_ptr,
562
eagle_input_ids_ptr,
563
eagle_positions_ptr,
564
+ eagle_query_start_loc_ptr,
565
+ eagle_seq_lens_ptr,
566
target_input_ids_ptr,
567
target_positions_ptr,
568
idx_mapping_ptr,
@@ -472,20 +571,24 @@ def _prepare_eagle_inputs_kernel(
571
num_sampled_ptr,
572
num_rejected_ptr,
573
query_start_loc_ptr,
574
+ seq_lens_ptr,
575
+ max_num_reqs,
576
BLOCK_SIZE: tl.constexpr,
577
):
477
- batch_idx = tl.program_id(0)
478
- req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
578
+ req_idx = tl.program_id(0)
579
+ num_reqs = tl.num_programs(0)
580
+ req_state_idx = tl.load(idx_mapping_ptr + req_idx)
581
480
- query_start = tl.load(query_start_loc_ptr + batch_idx)
481
- query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
582
+ query_start = tl.load(query_start_loc_ptr + req_idx)
583
+ query_end = tl.load(query_start_loc_ptr + req_idx + 1)
584
query_len = query_end - query_start
585
+ seq_len = tl.load(seq_lens_ptr + req_idx)
586
587
# Get the true query length and next token after accounting for rejected tokens.
485
- num_rejected = tl.load(num_rejected_ptr + batch_idx)
588
+ num_rejected = tl.load(num_rejected_ptr + req_idx)
589
query_len -= num_rejected
590
488
- num_sampled = tl.load(num_sampled_ptr + batch_idx)
591
+ num_sampled = tl.load(num_sampled_ptr + req_idx)
592
if num_sampled > 0:
593
next_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
594
else:
@@ -501,7 +604,7 @@ def _prepare_eagle_inputs_kernel(
604
tl.store(eagle_input_ids_ptr + query_start + block - 1, input_ids, mask=mask)
605
606
last_token_index = query_start + query_len - 1
504
- tl.store(last_token_indices_ptr + batch_idx, last_token_index)
607
+ tl.store(last_token_indices_ptr + req_idx, last_token_index)
608
tl.store(eagle_input_ids_ptr + last_token_index, next_token)
609
610
# Copy positions.
@@ -511,11 +614,29 @@ def _prepare_eagle_inputs_kernel(
614
target_pos = tl.load(target_positions_ptr + query_start + block, mask=mask)
615
tl.store(eagle_positions_ptr + query_start + block, target_pos, mask=mask)
616
617
+ # Copy query start locations.
618
+ tl.store(eagle_query_start_loc_ptr + req_idx, query_start)
619
+ # Copy sequence lengths.
620
+ tl.store(eagle_seq_lens_ptr + req_idx, seq_len)
621
+ if req_idx == (num_reqs - 1):
622
+ # Pad query_start_loc for CUDA graphs.
623
+ for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE):
624
+ block = i + tl.arange(0, BLOCK_SIZE)
625
+ mask = block < max_num_reqs + 1
626
+ tl.store(eagle_query_start_loc_ptr + block, query_end, mask=mask)
627
+ # Pad seq_lens for CUDA graphs.
628
+ for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
629
+ block = i + tl.arange(0, BLOCK_SIZE)
630
+ mask = block < max_num_reqs
631
+ tl.store(eagle_seq_lens_ptr + block, 0, mask=mask)
632
+
633
634
def prepare_eagle_inputs(
635
input_buffers: InputBuffers,
636
input_batch: InputBatch,
637
# [num_reqs]
638
+ last_token_indices: torch.Tensor,
639
+ # [num_reqs]
640
num_sampled: torch.Tensor,
641
# [num_reqs]
642
num_rejected: torch.Tensor,
@@ -523,17 +644,15 @@ def prepare_eagle_inputs(
644
last_sampled: torch.Tensor,
645
# [max_num_reqs]
646
next_prefill_tokens: torch.Tensor,
647
+ max_num_reqs,
648
) -> torch.Tensor:
649
num_reqs = input_batch.num_reqs
528
- last_token_indices = torch.empty(
529
- num_reqs,
530
- dtype=torch.int64,
531
- device=num_sampled.device,
532
- )
650
_prepare_eagle_inputs_kernel[(num_reqs,)](
651
last_token_indices,
652
input_buffers.input_ids,
653
input_buffers.positions,
654
+ input_buffers.query_start_loc,
655
+ input_buffers.seq_lens,
656
input_batch.input_ids,
657
input_batch.positions,
658
input_batch.idx_mapping,
@@ -542,6 +661,8 @@ def prepare_eagle_inputs(
661
num_sampled,
662
num_rejected,
663
input_batch.query_start_loc,
664
+ input_batch.seq_lens,
665
+ max_num_reqs,
666
BLOCK_SIZE=1024,
667
)
668
return last_token_indices
@@ -550,18 +671,13 @@ def prepare_eagle_inputs(
671
@triton.jit
672
def _prepare_eagle_docode_kernel(
673
draft_tokens_ptr,
553
- output_hidden_states_ptr,
554
- output_hidden_states_stride,
555
- last_token_indices_ptr,
674
+ draft_tokens_stride,
675
target_seq_lens_ptr,
676
num_rejected_ptr,
677
input_ids_ptr,
678
positions_ptr,
560
- input_hidden_states_ptr,
561
- input_hidden_states_stride,
679
query_start_loc_ptr,
680
seq_lens_ptr,
564
- hidden_size,
681
max_model_len,
682
max_num_reqs,
683
BLOCK_SIZE: tl.constexpr,
@@ -584,24 +700,9 @@ def _prepare_eagle_docode_kernel(
700
return
701
702
# draft token -> input id.
587
- draft_token = tl.load(draft_tokens_ptr + req_idx)
703
+ draft_token = tl.load(draft_tokens_ptr + req_idx * draft_tokens_stride)
704
tl.store(input_ids_ptr + req_idx, draft_token)
705
590
- # output hidden states -> input hidden states.
591
- src_idx = tl.load(last_token_indices_ptr + req_idx)
592
- for i in range(0, hidden_size, BLOCK_SIZE):
593
- block = i + tl.arange(0, BLOCK_SIZE)
594
- mask = block < hidden_size
595
- output_hidden_states = tl.load(
596
- output_hidden_states_ptr + src_idx * output_hidden_states_stride + block,
597
- mask=mask,
598
- )
599
- tl.store(
600
- input_hidden_states_ptr + req_idx * input_hidden_states_stride + block,
601
- output_hidden_states,
602
- mask=mask,
603
- )
604
-
706
# Compute position and seq_lens.
707
# NOTE(woosuk): To prevent out-of-range access, we clamp these values
708
# if they reach the max model length.
@@ -618,31 +719,22 @@ def _prepare_eagle_docode_kernel(
719
720
def prepare_eagle_decode(
721
draft_tokens: torch.Tensor,
621
- output_hidden_states: torch.Tensor,
622
- last_token_indices: torch.Tensor,
722
target_seq_lens: torch.Tensor,
723
num_rejected: torch.Tensor,
724
input_buffers: InputBuffers,
626
- input_hidden_states: torch.Tensor,
725
max_model_len: int,
726
max_num_reqs: int,
727
):
728
num_reqs = draft_tokens.shape[0]
631
- hidden_size = output_hidden_states.shape[-1]
729
_prepare_eagle_docode_kernel[(num_reqs + 1,)](
730
draft_tokens,
634
- output_hidden_states,
635
- output_hidden_states.stride(0),
636
- last_token_indices,
731
+ draft_tokens.stride(0),
732
target_seq_lens,
733
num_rejected,
734
input_buffers.input_ids,
735
input_buffers.positions,
641
- input_hidden_states,
642
- input_hidden_states.stride(0),
736
input_buffers.query_start_loc,
737
input_buffers.seq_lens,
645
- hidden_size,
738
max_model_len,
739
max_num_reqs,
740
BLOCK_SIZE=1024,