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 15, 13:30 UTC
[Bugfix] Get a specific type of layer from forward context (#17222)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Apr 27, 2025 at 15:58 UTC
838cedade77a10830a56ab58253dae1ed8fe6c81
5 files changed
+28
-23
vllm/attention/backends/flashinfer.py
+2
-4
@@ -38,7 +38,7 @@ from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,
38
is_block_tables_empty)
39
from vllm.attention.layer import Attention
40
from vllm.attention.ops.paged_attn import PagedAttention
41
-from vllm.config import VllmConfig
41
+from vllm.config import VllmConfig, get_layers_from_vllm_config
42
from vllm.logger import init_logger
43
from vllm.utils import (async_tensor_h2d, get_kv_cache_torch_dtype,
44
make_tensor_with_pad)
@@ -140,12 +140,10 @@ def get_per_layer_parameters(
140
to use during `plan`.
141
"""
142
143
- layers = vllm_config.compilation_config.static_forward_context
143
+ layers = get_layers_from_vllm_config(vllm_config, Attention)
144
per_layer_params: Dict[str, PerLayerParameters] = {}
145
146
for key, layer in layers.items():
147
- assert isinstance(layer, Attention)
148
-
147
impl = layer.impl
148
assert isinstance(impl, FlashInferImpl)
149
vllm/config.py
+15
-1
@@ -3445,7 +3445,8 @@ class CompilationConfig(BaseModel):
3445
compilation_time: float = PrivateAttr
3446
3447
# Per-model forward context
3448
- # Map from layer name to the attention cls
3448
+ # Map from layer name to layer objects that need to be accessed outside
3449
+ # model code, e.g., Attention, FusedMOE when dp_size>1.
3450
static_forward_context: dict[str, Any] = PrivateAttr
3451
3452
def compute_hash(self) -> str:
@@ -4079,3 +4080,16 @@ def assert_hashable(text):
4080
f"vLLM tried to hash some configs that may have Python objects ids "
4081
f"in them. This is a bug, please file an issue. "
4082
f"Text being hashed: {text}")
4083
+
4084
+
4085
+T = TypeVar("T")
4086
+
4087
+
4088
+def get_layers_from_vllm_config(vllm_config: VllmConfig,
4089
+ layer_type: type[T]) -> dict[str, T]:
4090
+ return {
4091
+ layer_name: layer
4092
+ for layer_name, layer in
4093
+ vllm_config.compilation_config.static_forward_context.items()
4094
+ if isinstance(layer, layer_type)
4095
+ }
vllm/v1/attention/backends/flashinfer.py
+3
-4
@@ -14,7 +14,8 @@ import vllm.envs as envs
14
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
15
AttentionType)
16
from vllm.attention.layer import Attention
17
-from vllm.config import VllmConfig, get_current_vllm_config
17
+from vllm.config import (VllmConfig, get_current_vllm_config,
18
+ get_layers_from_vllm_config)
19
from vllm.logger import init_logger
20
from vllm.v1.attention.backends.flash_attn import use_cascade_attention
21
@@ -81,12 +82,10 @@ def get_per_layer_parameters(
82
to use during `plan`.
83
"""
84
84
- layers = vllm_config.compilation_config.static_forward_context
85
+ layers = get_layers_from_vllm_config(vllm_config, Attention)
86
per_layer_params: dict[str, PerLayerParameters] = {}
87
88
for key, layer in layers.items():
88
- assert isinstance(layer, Attention)
89
-
89
impl = layer.impl
90
assert isinstance(impl, FlashInferImpl)
91
vllm/v1/worker/gpu_model_runner.py
+5
-10
@@ -12,13 +12,13 @@ import torch.nn as nn
12
13
from vllm.attention import AttentionType, get_attn_backend
14
from vllm.attention.layer import Attention
15
-from vllm.config import CompilationLevel, VllmConfig
15
+from vllm.config import (CompilationLevel, VllmConfig,
16
+ get_layers_from_vllm_config)
17
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
18
has_kv_transfer_group)
19
from vllm.distributed.parallel_state import get_pp_group, graph_capture
20
from vllm.forward_context import set_forward_context
21
from vllm.logger import init_logger
21
-from vllm.model_executor.layers.fused_moe import FusedMoE
22
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
23
from vllm.model_executor.model_loader import get_model
24
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -1733,17 +1733,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1733
format. Layers that do not need KV cache are not included.
1734
"""
1735
1736
- forward_ctx = self.vllm_config.compilation_config.static_forward_context
1736
+ layers = get_layers_from_vllm_config(self.vllm_config, Attention)
1737
block_size = self.vllm_config.cache_config.block_size
1738
use_mla = self.vllm_config.model_config.use_mla
1739
kv_cache_spec: dict[str, KVCacheSpec] = {}
1740
- for layer_name, attn_module in forward_ctx.items():
1741
- if isinstance(attn_module, FusedMoE):
1742
- continue
1743
-
1744
- # TODO: Support other attention modules, e.g., sliding window,
1745
- # cross-attention
1746
- assert isinstance(attn_module, Attention)
1740
+ for layer_name, attn_module in layers.items():
1741
+ # TODO: Support other attention modules, e.g., cross-attention
1742
if attn_module.attn_type == AttentionType.DECODER:
1743
if attn_module.sliding_window is not None:
1744
kv_cache_spec[layer_name] = SlidingWindowSpec(
vllm/v1/worker/tpu_model_runner.py
+3
-4
@@ -17,7 +17,7 @@ import vllm.envs as envs
17
from vllm.attention.backends.abstract import AttentionType
18
from vllm.attention.layer import Attention
19
from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
20
-from vllm.config import VllmConfig
20
+from vllm.config import VllmConfig, get_layers_from_vllm_config
21
from vllm.forward_context import set_forward_context
22
from vllm.logger import init_logger
23
from vllm.model_executor.model_loader import get_model
@@ -429,11 +429,10 @@ class TPUModelRunner:
429
format. Layers that do not need KV cache are not included.
430
"""
431
432
- forward_ctx = self.vllm_config.compilation_config.static_forward_context
432
+ layers = get_layers_from_vllm_config(self.vllm_config, Attention)
433
block_size = self.vllm_config.cache_config.block_size
434
kv_cache_spec: dict[str, KVCacheSpec] = {}
435
- for layer_name, attn_module in forward_ctx.items():
436
- assert isinstance(attn_module, Attention)
435
+ for layer_name, attn_module in layers.items():
436
if attn_module.attn_type == AttentionType.DECODER:
437
if attn_module.sliding_window is not None:
438
kv_cache_spec[layer_name] = SlidingWindowSpec(