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 2, 08:30 UTC
[MRV2] Remove Eagle's dedicated CUDA graph pool (#44078)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jun 1, 2026 at 01:00 UTC
4721bb3aa43078167eb893a9ebf9e50565030c1c
2 files changed
+2
-26
vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
+2
-23
@@ -4,7 +4,6 @@ from collections.abc import Callable
4
5
import torch
6
7
-from vllm.config import VllmConfig
7
from vllm.config.compilation import CUDAGraphMode
8
from vllm.v1.kv_cache_interface import KVCacheConfig
9
from vllm.v1.worker.gpu.block_table import BlockTables
@@ -19,27 +18,7 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState
18
from vllm.v1.worker.utils import AttentionGroup
19
20
22
-class EagleCudaGraphManagerBase(CudaGraphManager):
23
- """Base CudaGraphManager for Eagle with a dedicated graph pool."""
24
-
25
- def __init__(
26
- self,
27
- vllm_config: VllmConfig,
28
- device: torch.device,
29
- cudagraph_mode: CUDAGraphMode,
30
- decode_query_len: int,
31
- ):
32
- super().__init__(vllm_config, device, cudagraph_mode, decode_query_len)
33
-
34
- # Use a dedicated pool for Eagle to avoid memory overlap with the main
35
- # model's cudagraph. The base class uses a shared global pool, but Eagle's
36
- # internal allocations (e.g., gumbel_sample temporaries) can conflict with
37
- # the main model's allocations when sharing the same pool.
38
- if cudagraph_mode:
39
- self.pool = torch.cuda.graph_pool_handle()
40
-
41
-
42
-class PrefillEagleCudaGraphManager(EagleCudaGraphManagerBase):
21
+class PrefillEagleCudaGraphManager(CudaGraphManager):
22
"""Eagle CudaGraphManager for prefill, using pre-built attention states
23
from the target model's capture."""
24
@@ -74,7 +53,7 @@ class PrefillEagleCudaGraphManager(EagleCudaGraphManagerBase):
53
super().capture(create_forward_fn, progress_bar_desc)
54
55
77
-class DecodeEagleCudaGraphManager(EagleCudaGraphManagerBase):
56
+class DecodeEagleCudaGraphManager(CudaGraphManager):
57
"""Eagle CudaGraphManager for decode draft generation, building its own
58
attention metadata from scratch."""
59
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
-3
@@ -145,9 +145,6 @@ class EagleSpeculator:
145
cudagraph_mode,
146
decode_query_len=1,
147
)
148
- # Share a single pool between prefill and decode since they never
149
- # execute concurrently.
150
- self.decode_cudagraph_manager.pool = self.prefill_cudagraph_manager.pool
148
149
def load_model(self, target_model: nn.Module) -> None:
150
target_attn_layer_names = get_layers_from_vllm_config(