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, 04:45 UTC
[MRV2] Also enable MRV2 for Llama and Mistral dense models (#43458)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: yewentao256 <zhyanwentao@126.com>
Nick Hill committed
Jun 2, 2026 at 11:18 UTC
da107a59e56b13a170ab5751c93cd22491baf641
5 files changed
+59
-9
tests/test_config.py
+20
@@ -90,6 +90,26 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):
90
),
91
True,
92
),
93
+ (
94
+ SimpleNamespace(
95
+ model="meta-llama/Llama-3.2-1B",
96
+ architectures=["LlamaForCausalLM"],
97
+ runner_type="generate",
98
+ is_moe=False,
99
+ is_quantized=False,
100
+ ),
101
+ True,
102
+ ),
103
+ (
104
+ SimpleNamespace(
105
+ model="mistralai/Mistral-7B-v0.1",
106
+ architectures=["MistralForCausalLM"],
107
+ runner_type="generate",
108
+ is_moe=False,
109
+ is_quantized=False,
110
+ ),
111
+ True,
112
+ ),
113
(
114
SimpleNamespace(
115
model="facebook/opt-125m",
tests/v1/engine/test_abort_final_step.py
+25
-8
@@ -184,14 +184,31 @@ async def test_abort_during_final_step(async_scheduling: bool):
184
original_execute_model = Worker.execute_model
185
186
def execute_model_with_wait(self, scheduler_output):
187
- # Signal that execute_model has been called by deleting ready_file
188
- if ready_file.exists():
189
- ready_file.unlink()
190
-
191
- # Wait for the block file to be deleted (triggered from test after abort)
192
- # This runs in the worker process (after fork), so we poll the filesystem
193
- while block_file.exists():
194
- time.sleep(0.01)
187
+ # V2's `gpu_worker.compile_or_warm_up_model` calls
188
+ # `warmup_kernels(...)` during engine init, which itself calls
189
+ # `Worker.execute_model` three times (prefill / decode / cleanup)
190
+ # to JIT compile triton kernels. None of those carry the test's
191
+ # request id, so we only stall when our actual request is being
192
+ # processed.
193
+ scheduled = scheduler_output.num_scheduled_tokens or {}
194
+ finished = scheduler_output.finished_req_ids or set()
195
+
196
+ def is_target_request(req_ids):
197
+ return any(
198
+ rid == request_id or rid.startswith(f"{request_id}-")
199
+ for rid in req_ids
200
+ )
201
+
202
+ if is_target_request(scheduled) or is_target_request(finished):
203
+ # Signal that execute_model has been called by deleting ready_file
204
+ if ready_file.exists():
205
+ ready_file.unlink()
206
+
207
+ # Wait for the block file to be deleted (triggered from test after
208
+ # abort). This runs in the worker process (after fork), so we poll
209
+ # the filesystem.
210
+ while block_file.exists():
211
+ time.sleep(0.01)
212
return original_execute_model(self, scheduler_output)
213
214
# Patch execute_model to inject the wait
tests/v1/shutdown/test_forward_error.py
+1
@@ -36,6 +36,7 @@ def evil_forward(self, *args, **kwargs):
36
raise Exception("Simulated illegal memory access on Rank 0!")
37
self.num_calls += 1
38
39
+ kwargs.setdefault("intermediate_tensors", None) # required for MRV2
40
return self.model(*args, **kwargs)
41
42
vllm/config/vllm.py
+7
-1
@@ -66,7 +66,13 @@ else:
66
67
logger = init_logger(__name__)
68
69
-DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset({"Qwen3ForCausalLM"})
69
+DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
70
+ {
71
+ "LlamaForCausalLM",
72
+ "MistralForCausalLM",
73
+ "Qwen3ForCausalLM",
74
+ }
75
+)
76
77
78
class OptimizationLevel(IntEnum):
vllm/v1/worker/gpu/model_runner.py
+6
@@ -344,6 +344,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
344
self.reset_encoder_cache()
345
self.reset_mm_cache()
346
347
+ def apply_sparse_weight_patches(self, *args, **kwargs) -> None:
348
+ # TODO: Use full version instead of import when fully migrated to v2
349
+ from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1
350
+
351
+ GPUModelRunnerV1.apply_sparse_weight_patches(self, *args, **kwargs) # type: ignore[arg-type]
352
+
353
def update_config(self, *args, **kwargs) -> None:
354
# TODO(Wentao): Use full version instead of import when fully migrated to v2
355
from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1