diff --git a/tests/test_config.py b/tests/test_config.py index c0bd4b14ff8..b78570e54fb 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -90,6 +90,26 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): ), True, ), + ( + SimpleNamespace( + model="meta-llama/Llama-3.2-1B", + architectures=["LlamaForCausalLM"], + runner_type="generate", + is_moe=False, + is_quantized=False, + ), + True, + ), + ( + SimpleNamespace( + model="mistralai/Mistral-7B-v0.1", + architectures=["MistralForCausalLM"], + runner_type="generate", + is_moe=False, + is_quantized=False, + ), + True, + ), ( SimpleNamespace( model="facebook/opt-125m", diff --git a/tests/v1/engine/test_abort_final_step.py b/tests/v1/engine/test_abort_final_step.py index 8f1e8029955..d8d5b73d45b 100644 --- a/tests/v1/engine/test_abort_final_step.py +++ b/tests/v1/engine/test_abort_final_step.py @@ -184,14 +184,31 @@ async def test_abort_during_final_step(async_scheduling: bool): original_execute_model = Worker.execute_model def execute_model_with_wait(self, scheduler_output): - # Signal that execute_model has been called by deleting ready_file - if ready_file.exists(): - ready_file.unlink() + # V2's `gpu_worker.compile_or_warm_up_model` calls + # `warmup_kernels(...)` during engine init, which itself calls + # `Worker.execute_model` three times (prefill / decode / cleanup) + # to JIT compile triton kernels. None of those carry the test's + # request id, so we only stall when our actual request is being + # processed. + scheduled = scheduler_output.num_scheduled_tokens or {} + finished = scheduler_output.finished_req_ids or set() - # Wait for the block file to be deleted (triggered from test after abort) - # This runs in the worker process (after fork), so we poll the filesystem - while block_file.exists(): - time.sleep(0.01) + def is_target_request(req_ids): + return any( + rid == request_id or rid.startswith(f"{request_id}-") + for rid in req_ids + ) + + if is_target_request(scheduled) or is_target_request(finished): + # Signal that execute_model has been called by deleting ready_file + if ready_file.exists(): + ready_file.unlink() + + # Wait for the block file to be deleted (triggered from test after + # abort). This runs in the worker process (after fork), so we poll + # the filesystem. + while block_file.exists(): + time.sleep(0.01) return original_execute_model(self, scheduler_output) # Patch execute_model to inject the wait diff --git a/tests/v1/shutdown/test_forward_error.py b/tests/v1/shutdown/test_forward_error.py index eadb1abb6d5..8bc09a64cac 100644 --- a/tests/v1/shutdown/test_forward_error.py +++ b/tests/v1/shutdown/test_forward_error.py @@ -36,6 +36,7 @@ def evil_forward(self, *args, **kwargs): raise Exception("Simulated illegal memory access on Rank 0!") self.num_calls += 1 + kwargs.setdefault("intermediate_tensors", None) # required for MRV2 return self.model(*args, **kwargs) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index db139b6532a..b772a25028f 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -66,7 +66,13 @@ else: logger = init_logger(__name__) -DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset({"Qwen3ForCausalLM"}) +DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( + { + "LlamaForCausalLM", + "MistralForCausalLM", + "Qwen3ForCausalLM", + } +) class OptimizationLevel(IntEnum): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0f96e5c2d1d..532b3c6c0b8 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -344,6 +344,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.reset_encoder_cache() self.reset_mm_cache() + def apply_sparse_weight_patches(self, *args, **kwargs) -> None: + # TODO: Use full version instead of import when fully migrated to v2 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 + + GPUModelRunnerV1.apply_sparse_weight_patches(self, *args, **kwargs) # type: ignore[arg-type] + def update_config(self, *args, **kwargs) -> None: # TODO(Wentao): Use full version instead of import when fully migrated to v2 from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1