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, 12:00 UTC
[V0 Deprecation] Deprecate virtual engine (#37195)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Mar 18, 2026 at 17:30 UTC
0d81a1fe6190f47379c9905be5757e7b6bba5d14
23 files changed
+23
-45
tests/compile/passes/test_rope_kvcache_fusion.py
+2
-2
@@ -295,7 +295,7 @@ def test_rope_kvcache_fusion(
295
}
296
q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
297
attn_layer = forward_context.no_compile_layers[model.layer_name]
298
- kv_cache_unfused = attn_layer.kv_cache[forward_context.virtual_engine]
298
+ kv_cache_unfused = attn_layer.kv_cache[0]
299
del dummy
300
301
torch._dynamo.mark_dynamic(qkv, 0)
@@ -309,7 +309,7 @@ def test_rope_kvcache_fusion(
309
}
310
q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
311
attn_layer = forward_context.no_compile_layers[model.layer_name]
312
- kv_cache_fused = attn_layer.kv_cache[forward_context.virtual_engine]
312
+ kv_cache_fused = attn_layer.kv_cache[0]
313
del dummy
314
315
assert fusion_pass.matched_count == 1
tests/v1/kv_connector/unit/test_decode_bench_connector.py
+1
-1
@@ -86,7 +86,7 @@ class DecodeBenchTestRunner:
86
self._block_hasher = get_request_block_hasher(block_size, sha256)
87
88
self._dummy_ctx: ForwardContext = ForwardContext(
89
- no_compile_layers={}, attn_metadata={}, virtual_engine=0, slot_mapping={}
89
+ no_compile_layers={}, attn_metadata={}, slot_mapping={}
90
)
91
92
def new_request(self, token_ids: list[int]) -> Request:
tests/v1/kv_connector/unit/test_lmcache_integration.py
-1
@@ -211,7 +211,6 @@ def test_forward_context_interface():
211
from vllm.forward_context import ForwardContext
212
213
assumes(ForwardContext, "no_compile_layers", is_instance_of=dict)
214
- assumes(ForwardContext, "virtual_engine")
214
assumes(ForwardContext, "attn_metadata")
215
216
tests/v1/kv_connector/unit/test_nixl_connector.py
-8
@@ -599,7 +599,6 @@ class TestNixlHandshake:
599
dummy_ctx = ForwardContext(
600
no_compile_layers={},
601
attn_metadata={},
602
- virtual_engine=0,
602
slot_mapping={},
603
)
604
_before_load = time.perf_counter()
@@ -672,7 +671,6 @@ class TestNixlHandshake:
671
dummy_ctx = ForwardContext(
672
no_compile_layers={},
673
attn_metadata={},
675
- virtual_engine=0,
674
slot_mapping={},
675
)
676
_before_load = time.perf_counter()
@@ -908,7 +906,6 @@ class TestNixlHandshake:
906
dummy_ctx = ForwardContext(
907
no_compile_layers={},
908
attn_metadata={},
911
- virtual_engine=0,
909
slot_mapping={},
910
)
911
_before_load = time.perf_counter()
@@ -1079,7 +1076,6 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
1076
dummy_ctx = ForwardContext(
1077
no_compile_layers={},
1078
attn_metadata={},
1082
- virtual_engine=0,
1079
slot_mapping={},
1080
)
1081
connector.start_load_kv(dummy_ctx)
@@ -1890,7 +1886,6 @@ def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_
1886
dummy_ctx = ForwardContext(
1887
no_compile_layers={},
1888
attn_metadata={},
1893
- virtual_engine=0,
1889
slot_mapping={},
1890
)
1891
connector.start_load_kv(dummy_ctx)
@@ -2059,7 +2054,6 @@ def test_transfer_failure_logging(
2054
dummy_ctx = ForwardContext(
2055
no_compile_layers={},
2056
attn_metadata={},
2062
- virtual_engine=0,
2057
slot_mapping={},
2058
)
2059
@@ -2162,7 +2156,6 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
2156
dummy_ctx = ForwardContext(
2157
no_compile_layers={},
2158
attn_metadata={},
2165
- virtual_engine=0,
2159
slot_mapping={},
2160
)
2161
connector.start_load_kv(dummy_ctx)
@@ -2215,7 +2208,6 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
2208
dummy_ctx = ForwardContext(
2209
no_compile_layers={},
2210
attn_metadata={},
2218
- virtual_engine=0,
2211
slot_mapping={},
2212
)
2213
connector.start_load_kv(dummy_ctx)
tests/v1/kv_connector/unit/test_offloading_connector.py
-1
@@ -261,7 +261,6 @@ class RequestRunner:
261
self._dummy_ctx: ForwardContext = ForwardContext(
262
no_compile_layers={},
263
attn_metadata={},
264
- virtual_engine=0,
264
slot_mapping={},
265
)
266
vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
+1
-1
@@ -185,7 +185,7 @@ class ExampleConnector(KVConnectorBase_V1):
185
if kv_cache_attr is None:
186
continue
187
188
- kv_cache_layer = kv_cache_attr[forward_context.virtual_engine]
188
+ kv_cache_layer = kv_cache_attr[0]
189
190
filename = self._generate_filename_debug(
191
layer_name, request.token_ids, request.mm_hashes
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
+1
-3
@@ -778,9 +778,7 @@ class LMCacheConnectorV1Impl:
778
continue
779
780
if layer_name not in self.kv_caches:
781
- self.kv_caches[layer_name] = attn_layer.kv_cache[
782
- forward_context.virtual_engine
783
- ]
781
+ self.kv_caches[layer_name] = attn_layer.kv_cache[0]
782
783
####################
784
# Worker side APIs
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py
+1
-1
@@ -214,7 +214,7 @@ class P2pNcclConnector(KVConnectorBase_V1):
214
if kv_cache is None:
215
continue
216
217
- layer = kv_cache[forward_context.virtual_engine]
217
+ layer = kv_cache[0]
218
219
kv_cache = self.p2p_nccl_engine.recv_tensor(
220
request.request_id + "#" + layer_name, remote_address
vllm/forward_context.py
-7
@@ -197,8 +197,6 @@ class ForwardContext:
197
for each microbatch.
198
Set dynamically for each forward pass
199
"""
200
- # TODO: remove after making all virtual_engines share the same kv cache
201
- virtual_engine: int # set dynamically for each forward pass
200
# set dynamically for each forward pass
201
dp_metadata: DPMetadata | None = None
202
# determine the cudagraph style at runtime to be FULL, PIECEWISE, or NONE.
@@ -265,7 +263,6 @@ def is_forward_context_available() -> bool:
263
def create_forward_context(
264
attn_metadata: Any,
265
vllm_config: VllmConfig,
268
- virtual_engine: int = 0,
266
dp_metadata: DPMetadata | None = None,
267
cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
268
batch_descriptor: BatchDescriptor | None = None,
@@ -282,7 +279,6 @@ def create_forward_context(
279
return ForwardContext(
280
no_compile_layers=vllm_config.compilation_config.static_forward_context,
281
all_moe_layers=all_moe_layers,
285
- virtual_engine=virtual_engine,
282
attn_metadata=attn_metadata,
283
slot_mapping=slot_mapping or {},
284
dp_metadata=dp_metadata,
@@ -313,7 +309,6 @@ def override_forward_context(forward_context: ForwardContext | None):
309
def set_forward_context(
310
attn_metadata: Any,
311
vllm_config: VllmConfig,
316
- virtual_engine: int = 0,
312
num_tokens: int | None = None,
313
num_tokens_across_dp: torch.Tensor | None = None,
314
cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
@@ -362,7 +357,6 @@ def set_forward_context(
357
additional_kwargs = current_platform.set_additional_forward_context(
358
attn_metadata=attn_metadata,
359
vllm_config=vllm_config,
365
- virtual_engine=virtual_engine,
360
dp_metadata=dp_metadata,
361
num_tokens=num_tokens,
362
num_tokens_across_dp=num_tokens_across_dp,
@@ -374,7 +368,6 @@ def set_forward_context(
368
forward_context = create_forward_context(
369
attn_metadata,
370
vllm_config,
377
- virtual_engine,
371
dp_metadata,
372
cudagraph_runtime_mode,
373
batch_descriptor,
vllm/model_executor/layers/attention/attention.py
+2
-2
@@ -589,7 +589,7 @@ def get_attention_context(
589
- attn_metadata: Attention metadata for this specific layer, or None if
590
no metadata available
591
- attn_layer: The attention layer instance (Attention or MLAAttention)
592
- - kv_cache: The KV cache tensor for current virtual engine
592
+ - kv_cache: The KV cache tensor for current forward pass
593
- slot_mapping: The slot mapping for this specific layer
594
595
Note: attn_metadata may be None, but attn_layer and kv_cache are always
@@ -600,7 +600,7 @@ def get_attention_context(
600
if isinstance(attn_metadata, dict):
601
attn_metadata = attn_metadata[layer_name]
602
attn_layer: Attention | MLAAttention = forward_context.no_compile_layers[layer_name]
603
- kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
603
+ kv_cache = attn_layer.kv_cache[0]
604
slot_mapping = forward_context.slot_mapping
605
assert isinstance(slot_mapping, dict), (
606
f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
vllm/model_executor/layers/attention/mla_attention.py
+2
-2
@@ -480,7 +480,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
480
attn_metadata = forward_context.attn_metadata
481
if isinstance(attn_metadata, dict):
482
attn_metadata = attn_metadata[self.layer_name]
483
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
483
+ self_kv_cache = self.kv_cache[0]
484
slot_mapping = forward_context.slot_mapping
485
486
assert isinstance(slot_mapping, dict), (
@@ -940,7 +940,7 @@ def unified_mla_kv_cache_update(
940
return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
941
942
attn_layer = forward_context.no_compile_layers[layer_name]
943
- kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
943
+ kv_cache = attn_layer.kv_cache[0]
944
945
slot_mapping = forward_context.slot_mapping
946
assert isinstance(slot_mapping, dict), (
vllm/model_executor/layers/attention/static_sink_attention.py
+1
-2
@@ -168,8 +168,7 @@ class StaticSinkAttention(Attention, CustomOp):
168
"sink_key and sink_value have not been prepared"
169
)
170
if not self.sink_populated:
171
- forward_context: ForwardContext = get_forward_context()
172
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
171
+ self_kv_cache = self.kv_cache[0]
172
torch.ops.vllm.maybe_populate_sink(self_kv_cache, self.layer_name)
173
174
return super().forward(query, key, value, output_shape)
vllm/model_executor/layers/kda.py
+1
-1
@@ -306,7 +306,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):
306
non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
307
non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
308
num_actual_tokens = attn_metadata.num_actual_tokens
309
- constant_caches = self.kv_cache[forward_context.virtual_engine]
309
+ constant_caches = self.kv_cache[0]
310
311
q_proj_states = q_proj_states[:num_actual_tokens]
312
k_proj_states = k_proj_states[:num_actual_tokens]
vllm/model_executor/layers/mamba/linear_attn.py
+1
-1
@@ -413,7 +413,7 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
413
qkvact = qkvact.view((qkv.shape[0], self.tp_heads, -1))
414
q, k, v = torch.split(qkvact, [self.head_dim] * 3, dim=-1)
415
if attn_metadata is not None:
416
- kv_cache = self.kv_cache[forward_context.virtual_engine][0]
416
+ kv_cache = self.kv_cache[0][0]
417
state_indices_tensor = attn_metadata.state_indices_tensor
418
clear_linear_attention_cache_for_new_sequences(
419
kv_cache, state_indices_tensor, attn_metadata
vllm/model_executor/layers/mamba/mamba_mixer.py
+1
-1
@@ -267,7 +267,7 @@ class MambaMixer(MambaBase, PluggableLayer):
267
query_start_loc_p = attn_metadata.query_start_loc_p
268
state_indices_tensor_p = attn_metadata.state_indices_tensor_p
269
state_indices_tensor_d = attn_metadata.state_indices_tensor_d
270
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
270
+ self_kv_cache = self.kv_cache[0]
271
conv_state = self_kv_cache[0].transpose(-1, -2)
272
ssm_state = self_kv_cache[1]
273
has_initial_states_p = attn_metadata.has_initial_states_p
vllm/model_executor/layers/mamba/mamba_mixer2.py
+1
-1
@@ -575,7 +575,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
575
assert isinstance(attn_metadata, dict)
576
attn_metadata = attn_metadata[self.prefix]
577
assert isinstance(attn_metadata, Mamba2AttentionMetadata)
578
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
578
+ self_kv_cache = self.kv_cache[0]
579
# conv_state = (..., dim, width-1) yet contiguous along 'dim'
580
conv_state = self_kv_cache[0].transpose(-1, -2)
581
ssm_state = self_kv_cache[1]
vllm/model_executor/layers/mamba/short_conv.py
+1
-1
@@ -117,7 +117,7 @@ class ShortConv(MambaBase, CustomOp):
117
assert isinstance(attn_metadata, dict)
118
attn_metadata = attn_metadata[self.prefix]
119
assert isinstance(attn_metadata, ShortConvAttentionMetadata)
120
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
120
+ self_kv_cache = self.kv_cache[0]
121
conv_state = self_kv_cache[0].transpose(-1, -2)
122
state_indices_tensor_p = attn_metadata.state_indices_tensor_p
123
state_indices_tensor_d = attn_metadata.state_indices_tensor_d
vllm/model_executor/models/bailing_moe_linear.py
+1
-1
@@ -709,7 +709,7 @@ class BailingMoELinearAttention(nn.Module, MambaBase):
709
710
# Get KV cache and state indices
711
if attn_metadata is not None:
712
- kv_cache = self.kv_cache[forward_context.virtual_engine][0]
712
+ kv_cache = self.kv_cache[0][0]
713
state_indices_tensor = attn_metadata.state_indices_tensor
714
clear_linear_attention_cache_for_new_sequences(
715
kv_cache, state_indices_tensor, attn_metadata
vllm/model_executor/models/extract_hidden_states.py
+1
-1
@@ -51,7 +51,7 @@ def unified_kv_cache_update(
51
"""
52
forward_context = get_forward_context()
53
attn_layer = forward_context.no_compile_layers[layer_name]
54
- kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
54
+ kv_cache = attn_layer.kv_cache[0]
55
56
slot_mapping = forward_context.slot_mapping
57
assert isinstance(slot_mapping, dict), (
vllm/model_executor/models/olmo_hybrid.py
+1
-1
@@ -428,7 +428,7 @@ class OlmoHybridGatedDeltaNet(nn.Module, MambaBase):
428
non_spec_token_indx = attn_metadata.non_spec_token_indx
429
spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor
430
non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor
431
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
431
+ self_kv_cache = self.kv_cache[0]
432
conv_state = self_kv_cache[0].transpose(-1, -2)
433
ssm_state = self_kv_cache[1]
434
num_actual_tokens = attn_metadata.num_actual_tokens
vllm/model_executor/models/plamo2.py
+1
-1
@@ -262,7 +262,7 @@ class Plamo2MambaMixer(MambaBase, PluggableLayer):
262
assert isinstance(attn_metadata, dict)
263
attn_metadata = attn_metadata[self.prefix]
264
assert isinstance(attn_metadata, Mamba2AttentionMetadata)
265
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
265
+ self_kv_cache = self.kv_cache[0]
266
# conv_state = (..., dim, width-1) yet contiguous along 'dim'
267
conv_state = self_kv_cache[0].transpose(-1, -2)
268
ssm_state = self_kv_cache[1]
vllm/model_executor/models/qwen3_next.py
+2
-4
@@ -842,7 +842,6 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
842
a=a,
843
core_attn_out=core_attn_out,
844
attn_metadata=attn_metadata,
845
- virtual_engine=forward_context.virtual_engine,
845
)
846
847
has_initial_state = attn_metadata.has_initial_state
@@ -853,7 +852,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
852
non_spec_token_indx = attn_metadata.non_spec_token_indx
853
spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501
854
non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
856
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
855
+ self_kv_cache = self.kv_cache[0]
856
conv_state = self_kv_cache[0].transpose(-1, -2)
857
ssm_state = self_kv_cache[1]
858
num_actual_tokens = attn_metadata.num_actual_tokens
@@ -1036,13 +1035,12 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
1035
a: torch.Tensor,
1036
core_attn_out: torch.Tensor,
1037
attn_metadata: GDNAttentionMetadata,
1039
- virtual_engine: int,
1038
):
1039
"""
1040
Core attention computation with a packed non-spec decode fast path.
1041
"""
1042
non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
1045
- self_kv_cache = self.kv_cache[virtual_engine]
1043
+ self_kv_cache = self.kv_cache[0]
1044
conv_state = self_kv_cache[0].transpose(-1, -2)
1045
ssm_state = self_kv_cache[1]
1046
num_actual_tokens = attn_metadata.num_actual_tokens
vllm/v1/worker/utils.py
+1
-1
@@ -510,7 +510,7 @@ def bind_kv_cache(
510
511
# Bind kv_caches to forward context
512
for layer_name, kv_cache in kv_caches.items():
513
- # NOTE: Use list because of v0 PP virtual engine.
513
+ # NOTE: Keep list wrapper for layers that index kv_cache by engine slot.
514
forward_context[layer_name].kv_cache = [kv_cache]
515
516