fix: match vLLM 0.31 generation profiling API
Accept the upstream randomize_inputs profiling flag while preserving structured codec inputs during generation-stage dummy runs. Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
1 parent
2974f8be59
commit
16f73d0da2
2 files changed
+20
-3
No files matched your search
@@ -3,6 +3,7 @@
|
||||
|
||||
import contextlib
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -22,6 +23,23 @@ class _StopExecutionError(Exception):
|
||||
"""Halts execute_model at the model boundary."""
|
||||
|
||||
|
||||
@pytest.mark.parametrize("randomize_inputs", [False, True])
|
||||
def test_generation_profile_accepts_upstream_randomize_inputs(randomize_inputs):
|
||||
"""Release workers can pass the profiling flag to codec-stage runners."""
|
||||
runner = object.__new__(GPUGenerationModelRunner)
|
||||
runner.supports_mm_inputs = False
|
||||
runner.max_num_tokens = 128
|
||||
runner.encoder_cache = {"temporary": object()}
|
||||
runner._dummy_run = Mock(return_value=(torch.zeros(1), None))
|
||||
runner._sync_device = Mock()
|
||||
|
||||
runner.profile_run(randomize_inputs=randomize_inputs)
|
||||
|
||||
runner._dummy_run.assert_called_once_with(128, is_profile=True, randomize_inputs=randomize_inputs)
|
||||
runner._sync_device.assert_called_once_with()
|
||||
assert runner.encoder_cache == {}
|
||||
|
||||
|
||||
class _DummyInputBatch:
|
||||
def __init__(self, num_reqs: int = 1):
|
||||
self.req_ids = [f"req-{i + 1}" for i in range(num_reqs)]
|
||||
|
||||
Reference in new issue
Block a user