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 6, 09:30 UTC
[Model Runner V2] fix draft attention metadata generation (#37364)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Mar 19, 2026 at 21:05 UTC
39474513f6631b1bc39a2400126bd7ff9394a774
3 files changed
+28
-6
vllm/v1/worker/gpu/attn_utils.py
+6
-1
@@ -30,7 +30,10 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
30
31
32
def init_attn_backend(
33
- kv_cache_config: KVCacheConfig, vllm_config: VllmConfig, device: torch.device
33
+ kv_cache_config: KVCacheConfig,
34
+ vllm_config: VllmConfig,
35
+ device: torch.device,
36
+ active_layer_names: set[str] | None = None,
37
):
38
attn_backends: dict[str, type[AttentionBackend]] = {}
39
attn_groups: list[list[AttentionGroup]] = []
@@ -39,6 +42,8 @@ def init_attn_backend(
42
kv_cache_config.kv_cache_groups
43
):
44
layer_names = kv_cache_group_spec.layer_names
45
+ if active_layer_names is not None:
46
+ layer_names = list(active_layer_names.intersection(layer_names))
47
48
layer_type = cast(type[Any], AttentionLayerBase)
49
attn_layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)
vllm/v1/worker/gpu/model_runner.py
-1
@@ -350,7 +350,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
350
self.speculator.set_attn(
351
self.model_state,
352
self.kv_cache_config,
353
- self.attn_groups,
353
self.block_tables,
354
)
355
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+22
-4
@@ -5,15 +5,17 @@ from typing import Any
5
import torch
6
import torch.nn as nn
7
8
-from vllm.config import VllmConfig
8
+from vllm.config import VllmConfig, get_layers_from_vllm_config
9
from vllm.config.compilation import CUDAGraphMode
10
from vllm.forward_context import BatchDescriptor, set_forward_context
11
from vllm.logger import init_logger
12
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
13
from vllm.triton_utils import tl, triton
14
from vllm.v1.kv_cache_interface import KVCacheConfig
15
from vllm.v1.worker.gpu.attn_utils import (
16
build_attn_metadata,
17
build_slot_mappings_by_layer,
18
+ init_attn_backend,
19
)
20
from vllm.v1.worker.gpu.block_table import BlockTables
21
from vllm.v1.worker.gpu.dp_utils import sync_cudagraph_and_dp_padding
@@ -22,7 +24,6 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState
24
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
25
from vllm.v1.worker.gpu.spec_decode.eagle.cudagraph import EagleCudaGraphManager
26
from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
25
-from vllm.v1.worker.utils import AttentionGroup
27
28
logger = init_logger(__name__)
29
@@ -87,18 +88,35 @@ class EagleSpeculator:
88
)
89
90
def load_model(self, target_model: nn.Module) -> None:
91
+ target_attn_layer_names = get_layers_from_vllm_config(
92
+ self.vllm_config,
93
+ AttentionLayerBase, # type: ignore[type-abstract]
94
+ ).keys()
95
+
96
self.model = load_eagle_model(target_model, self.vllm_config)
97
98
+ all_attn_layers = get_layers_from_vllm_config(
99
+ self.vllm_config,
100
+ AttentionLayerBase, # type: ignore[type-abstract]
101
+ ).keys()
102
+ self.draft_attn_layer_names = set(all_attn_layers) - set(
103
+ target_attn_layer_names
104
+ )
105
+
106
def set_attn(
107
self,
108
model_state: ModelState,
109
kv_cache_config: KVCacheConfig,
96
- attn_groups: list[list[AttentionGroup]],
110
block_tables: BlockTables,
111
) -> None:
112
self.model_state = model_state
113
self.kv_cache_config = kv_cache_config
101
- self.attn_groups = attn_groups
114
+ _, self.attn_groups = init_attn_backend(
115
+ kv_cache_config,
116
+ self.vllm_config,
117
+ self.device,
118
+ active_layer_names=self.draft_attn_layer_names,
119
+ )
120
self.block_tables = block_tables
121
122
@torch.inference_mode()