biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[WIP] TPU V1 Support Refactored (#13049)

Browse Source
This commit is contained in:
Alexander Matveev
2025-02-14 03:21:53 -05:00
committed by GitHub
parent b0ccfc565a
commit 45f90bcbba
8 changed files with 1738 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
tests/entrypoints/llm/test_accuracy.py
View File

@@ -21,10 +21,13 @@ RTOL = 0.03
EXPECTED_VALUE = 0.58
def run_test():
def run_test(more_args=None):
"""Run the end to end accuracy test."""
model_args = f"pretrained={MODEL_NAME},max_model_len=2048"
model_args = f"pretrained={MODEL_NAME},max_model_len=4096"
if more_args is not None:
model_args = "{},{}".format(model_args, more_args)
results = lm_eval.simple_evaluate(
model="vllm",
@@ -39,14 +42,21 @@ def run_test():
), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
@pytest.mark.skipif(not current_platform.is_cuda(),
reason="V1 is currently only supported on CUDA.")
@pytest.mark.skipif(not current_platform.is_cuda()
and not current_platform.is_tpu(),
reason="V1 is currently only supported on CUDA and TPU")
def test_lm_eval_accuracy_v1_engine(monkeypatch):
"""Run with the V1 Engine."""
with monkeypatch.context() as m:
m.setenv("VLLM_USE_V1", "1")
run_test()
more_args = None
if current_platform.is_tpu():
# Limit compilation time for TPU V1
more_args = "max_num_seqs=64"
run_test(more_args)
def test_lm_eval_accuracy_v0_engine(monkeypatch):