[CI] Organize speculative decoding E2E tests by coverage (#50330)

Signed-off-by: mgoin <mgoin64@gmail.com>
Signed-off-by: Michael Goin <mgoin64@gmail.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Michael Goin
2026-07-31 18:05:42 -07:00
committed by GitHub
co-authored by OpenAI Codex mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent b40d859c7b
commit fcdc7c2e9c
39 changed files with 1941 additions and 1863 deletions
+4 -2
View File
@@ -95,9 +95,11 @@ steps:
- vllm/triton_utils/
- vllm/utils/
- vllm/v1/
- tests/v1/e2e/spec_decode
- tests/v1/e2e/spec_decode/
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"'
pytest -v -s
v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism
v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism'
@@ -107,7 +107,6 @@ steps:
- vllm/v1/worker/gpu_worker.py
- tests/v1/spec_decode/test_max_len.py
- tests/v1/spec_decode/test_rejection_sampler_utils.py
- tests/v1/e2e/spec_decode/test_spec_decode.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+16 -23
View File
@@ -330,7 +330,7 @@ steps:
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- pytest -v -s v1/e2e/spec_decode/draft_model/
- label: Spec Decode Speculators + MTP # TBD
timeout_in_minutes: 180
@@ -348,7 +348,8 @@ steps:
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- pytest -v -s v1/e2e/spec_decode/speculators/
- pytest -v -s v1/e2e/spec_decode/mtp/
- label: V1 others (CPU) # TBD
timeout_in_minutes: 180
@@ -1836,7 +1837,7 @@ steps:
- tests/v1/spec_decode/test_max_len.py
- tests/v1/spec_decode/test_rejection_sampler_utils.py
- tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
- tests/v1/e2e/spec_decode/test_spec_decode.py
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
commands:
- set -x
@@ -1844,7 +1845,9 @@ steps:
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
- pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
- pytest -v -s v1/e2e/spec_decode/eagle/
- pytest -v -s v1/e2e/spec_decode/speculators/
- pytest -v -s v1/e2e/spec_decode/mtp/
#------------------------------------------------------ mi300 · models / basic -------------------------------------------------------#
@@ -2551,7 +2554,7 @@ steps:
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- pytest -v -s v1/e2e/spec_decode/draft_model/
- label: Spec Decode Eagle # TBD
timeout_in_minutes: 180
@@ -2569,7 +2572,7 @@ steps:
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
- pytest -v -s v1/e2e/spec_decode/eagle/
- label: Spec Decode Ngram + Suffix # TBD
timeout_in_minutes: 180
@@ -2587,7 +2590,7 @@ steps:
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
- pytest -v -s v1/e2e/spec_decode/ngram_suffix/
- label: Spec Decode Speculators + MTP # TBD
timeout_in_minutes: 180
@@ -2606,7 +2609,8 @@ steps:
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- pytest -v -s v1/e2e/spec_decode/speculators/
- pytest -v -s v1/e2e/spec_decode/mtp/
- label: Speculators Correctness # TBD
timeout_in_minutes: 180
@@ -2807,7 +2811,10 @@ steps:
- vllm/
- tests/v1/e2e
commands:
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
- >-
pytest -v -s
v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism
v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD
timeout_in_minutes: 180
@@ -2969,20 +2976,6 @@ steps:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
- label: V1 e2e (4 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_4
num_gpus: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/v1/e2e
commands:
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy"
- label: V1 e2e (4xH100-4xMI300) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+5 -30
View File
@@ -112,10 +112,13 @@ steps:
- vllm/triton_utils/
- vllm/utils/
- vllm/v1/
- tests/v1/e2e/spec_decode
- tests/v1/e2e/spec_decode/
commands:
# Only run tests that need exactly 2 GPUs
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
- >-
pytest -v -s
v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism
v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism
mirror:
amd:
dind: false
@@ -124,34 +127,6 @@ steps:
depends_on:
- image-build-amd
- label: V1 e2e (4 GPUs)
key: v1-e2e-4-gpus
timeout_in_minutes: 20 # TODO: Fix timeout after we have more confidence in the test stability
optional: true
num_devices: 4
source_file_dependencies:
- vllm/compilation/
- vllm/config/
- vllm/distributed/
- vllm/engine/
- vllm/envs.py
- vllm/forward_context.py
- vllm/inputs/
- vllm/logger.py
- vllm/logging_utils/
- vllm/model_executor/
- vllm/multimodal/
- vllm/platforms/
- vllm/sampling_params.py
- vllm/transformers_utils/
- vllm/triton_utils/
- vllm/utils/
- vllm/v1/
- tests/v1/e2e/spec_decode
commands:
# Only run tests that need 4 GPUs
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy"
- label: V1 e2e (4xH100)
key: v1-e2e-4xh100
timeout_in_minutes: 35
-27
View File
@@ -2,33 +2,6 @@ group: Miscellaneous
depends_on:
- image-build
steps:
- label: V1 Spec Decode
device: h200_35gb
key: v1-spec-decode
timeout_in_minutes: 40
source_file_dependencies:
- vllm/config/
- vllm/distributed/
- vllm/inputs/
- vllm/model_executor/
- vllm/platforms/
- vllm/sampling_params.py
- vllm/transformers_utils/
- vllm/utils/
- vllm/v1/
- tests/v1/spec_decode
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# TODO: create another `optional` test group for slow tests
- pytest -v -s -m 'not slow_test' v1/spec_decode
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 50
depends_on:
- image-build-amd
- label: V1 Sample + Logits
key: v1-sample-logits
timeout_in_minutes: 83
+4 -3
View File
@@ -103,18 +103,19 @@ steps:
- label: Model Runner V2 Spec Decode
device: h200_35gb
key: model-runner-v2-spec-decode
timeout_in_minutes: 30
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- tests/v1/spec_decode/test_max_len.py
- tests/v1/spec_decode/test_rejection_sampler_utils.py
- tests/v1/e2e/spec_decode/test_spec_decode.py
- tests/v1/e2e/spec_decode/
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
- pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
- pytest -v -s v1/e2e/spec_decode/eagle/
- pytest -v -s v1/e2e/spec_decode/speculators/
+78 -32
View File
@@ -2,16 +2,43 @@ group: Spec Decode
depends_on:
- image-build
steps:
- label: V1 Spec Decode
device: h200_35gb
key: v1-spec-decode
timeout_in_minutes: 40
source_file_dependencies:
- vllm/config/
- vllm/distributed/
- vllm/inputs/
- vllm/model_executor/
- vllm/platforms/
- vllm/sampling_params.py
- vllm/transformers_utils/
- vllm/utils/
- vllm/v1/
- tests/v1/spec_decode
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# TODO: create another `optional` test group for slow tests
- pytest -v -s -m 'not slow_test' v1/spec_decode
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 50
depends_on:
- image-build-amd
- label: Spec Decode Eagle
key: spec-decode-eagle
timeout_in_minutes: 25
device: h200_18gb
device: h200_35gb
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
- pytest -v -s v1/e2e/spec_decode/eagle/
mirror:
amd:
dind: false
@@ -28,7 +55,7 @@ steps:
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
- label: Spec Decode Eagle Nightly B200
- label: Spec Decode Eagle Nightly (B200)
key: spec-decode-eagle-nightly-b200
timeout_in_minutes: 25
device: b200-k8s
@@ -38,12 +65,12 @@ steps:
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
- pytest -v -s v1/e2e/spec_decode/eagle/
- label: Spec Decode Speculators + MTP
key: spec-decode-speculators-mtp
timeout_in_minutes: 20
device: h200_18gb
timeout_in_minutes: 50
device: h200_35gb
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
@@ -51,7 +78,8 @@ steps:
- vllm/transformers_utils/configs/speculators/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- pytest -v -s v1/e2e/spec_decode/speculators/
- pytest -v -s v1/e2e/spec_decode/mtp/
mirror:
amd:
dind: false
@@ -69,7 +97,7 @@ steps:
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
- label: Spec Decode Speculators + MTP Nightly B200
- label: Spec Decode Speculators + MTP Nightly (B200)
key: spec-decode-speculators-mtp-nightly-b200
timeout_in_minutes: 30
device: b200-k8s
@@ -80,19 +108,20 @@ steps:
- vllm/transformers_utils/configs/speculators/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- pytest -v -s v1/e2e/spec_decode/speculators/
- pytest -v -s v1/e2e/spec_decode/mtp/
- label: Spec Decode Ngram + Suffix
key: spec-decode-ngram-suffix
timeout_in_minutes: 20
device: h200_18gb
device: h200_35gb
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
- tests/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
- pytest -v -s v1/e2e/spec_decode/ngram_suffix/
- python3 spec_decode/test_custom_proposer.py
mirror:
amd:
@@ -119,7 +148,7 @@ steps:
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- pytest -v -s v1/e2e/spec_decode/draft_model/
mirror:
amd:
dind: false
@@ -136,7 +165,7 @@ steps:
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
- label: Spec Decode Draft Model Nightly B200
- label: Spec Decode Draft Model Nightly (B200)
key: spec-decode-draft-model-nightly-b200
timeout_in_minutes: 40
device: b200-k8s
@@ -146,9 +175,9 @@ steps:
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- pytest -v -s v1/e2e/spec_decode/draft_model/
- label: Speculators Correctness
- label: Speculators Correctness Nightly (H100)
key: speculators-correctness
timeout_in_minutes: 30
device: h100
@@ -156,25 +185,13 @@ steps:
num_devices: 1
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/dflash/
- vllm/model_executor/models/qwen3_dflash.py
- tests/v1/spec_decode/test_speculators_correctness.py
commands:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
- label: Spec Decode MTP hybrid (B200)
timeout_in_minutes: 20
device: b200-k8s
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/models/qwen3_5.py
- vllm/model_executor/models/qwen3_5_mtp.py
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
key: spec-decode-deepseek-mtp-parallel-load-b200
timeout_in_minutes: 30
@@ -191,14 +208,43 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
- label: Spec Decode Acceptance Rates Nightly
key: spec-decode-acceptance-rates-nightly
- label: Spec Decode AL DFlash Nightly
key: spec-decode-dflash-nightly
timeout_in_minutes: 90
device: h200_35gb
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/models/qwen3_dflash.py
- vllm/model_executor/models/laguna_dflash.py
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/
- label: Spec Decode AL DSpark Nightly
key: spec-decode-dspark-nightly
timeout_in_minutes: 60
device: h200_35gb
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/models/qwen3_dspark.py
- vllm/model_executor/models/gemma4_dspark.py
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "acceptance_rates"
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/
- label: Spec Decode AL MTP + Other Acceptance Nightly
key: spec-decode-mtp-other-acceptance-nightly
timeout_in_minutes: 60
device: h200_35gb
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/conftest.py
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/
+2 -2
View File
@@ -1449,8 +1449,8 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
max_num_seqs=32,
),
"DFlashLagunaForCausalLM": _HfExamplesInfo(
"poolside/Laguna-XS-2.1",
speculative_model="poolside/Laguna-XS-2.1-DFlash",
"poolside/Laguna-XS-2.1-NVFP4",
speculative_model="poolside/Laguna-XS-2.1-DFlash-NVFP4",
use_original_num_layers=True,
max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env
max_num_seqs=32,
+4 -1
View File
@@ -1983,7 +1983,10 @@ def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
return pytest.mark.skipif(
memory_gb < min_gb,
reason=f"Need at least {min_gb}GB GPU memory to run the test.",
reason=(
f"Need at least {min_gb}GB GPU memory to run the test "
f"(found {memory_gb}GB)."
),
)
@@ -18,7 +18,7 @@ SEED = 42
@pytest.fixture
def test_prompts():
"""
Adapted from tests/v1/e2e/spec_decode/test_spec_decode.py
Adapted from tests/v1/e2e/spec_decode/utils.py
"""
prompt_types = ["repeat", "sentence"]
# Setting higher num prompts increases the chance of numerics mismatch
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,152 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
import pytest
import torch
from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline
from tests.utils import single_gpu_only
from vllm import LLM
from vllm.distributed import cleanup_dist_env_and_memory
from ...utils import compute_acceptance_len
from ..utils import run_acceptance_length_eval
@dataclass(frozen=True)
class DFlashCorrectnessConfig:
model: str
draft_model: str
expected_accuracy: float
expected_acceptance_len: float
num_speculative_tokens: int = 16
max_model_len: int = 4096
max_num_seqs: int = 128
num_questions: int = 1319
use_chat_completions: bool = False
enforce_eager: bool = False
disable_flashinfer_sampler: bool = False
QWEN3_DFLASH = DFlashCorrectnessConfig(
model="Qwen/Qwen3-8B",
draft_model="z-lab/Qwen3-8B-DFlash-b16",
expected_accuracy=0.8,
expected_acceptance_len=3.5,
)
LAGUNA_DFLASH_NVFP4 = DFlashCorrectnessConfig(
model="poolside/Laguna-XS-2.1-NVFP4",
draft_model="poolside/Laguna-XS-2.1-DFlash-NVFP4",
expected_accuracy=0.7, # Standard GSM8K sanity floor.
expected_acceptance_len=3.55 * 0.9,
num_speculative_tokens=15,
max_model_len=8192,
max_num_seqs=32,
num_questions=200,
use_chat_completions=True,
enforce_eager=True,
disable_flashinfer_sampler=True,
)
@pytest.mark.parametrize("use_mrv2", [False, True])
def test_dflash_reference_acceptance_lengths(
monkeypatch: pytest.MonkeyPatch,
use_mrv2: bool,
):
run_acceptance_length_eval(
monkeypatch,
spec_config={
"model": QWEN3_DFLASH.model,
"trust_remote_code": True,
"speculative_config": {
"method": "dflash",
"model": QWEN3_DFLASH.draft_model,
"num_speculative_tokens": QWEN3_DFLASH.num_speculative_tokens,
"max_model_len": 32768,
},
"max_model_len": 32768,
"max_num_seqs": 128,
"gpu_memory_utilization": 0.85,
"enforce_eager": False,
"disable_log_stats": False,
},
# Table 1 in https://arxiv.org/pdf/2602.06036.
expected_acceptance_lengths={
"mt-bench": 4.24,
"humaneval": 6.50,
"gsm8k": 6.54 * 0.975,
},
chat_template_kwargs={"enable_thinking": False},
use_mrv2=use_mrv2,
)
@single_gpu_only
@pytest.mark.parametrize(
("config", "use_mrv2"),
[
pytest.param(
QWEN3_DFLASH,
False,
id="qwen3-mrv1",
),
pytest.param(
QWEN3_DFLASH,
True,
id="qwen3-mrv2",
),
pytest.param(
LAGUNA_DFLASH_NVFP4,
True,
id="laguna-nvfp4-mrv2",
),
],
)
def test_dflash_correctness(
monkeypatch: pytest.MonkeyPatch,
config: DFlashCorrectnessConfig,
use_mrv2: bool,
):
"""Guard GSM8K accuracy and batched acceptance length for DFlash models."""
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
if config.disable_flashinfer_sampler:
monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
spec_llm = LLM(
model=config.model,
trust_remote_code=True,
speculative_config={
"method": "dflash",
"model": config.draft_model,
"num_speculative_tokens": config.num_speculative_tokens,
"max_model_len": config.max_model_len,
},
max_model_len=config.max_model_len,
max_num_seqs=config.max_num_seqs,
gpu_memory_utilization=0.85,
enforce_eager=config.enforce_eager,
disable_log_stats=False,
)
results = evaluate_gsm8k_offline(
spec_llm,
num_questions=config.num_questions,
use_chat_completions=config.use_chat_completions,
)
accuracy = results["accuracy"]
acceptance_len = compute_acceptance_len(spec_llm.get_metrics())
print(
f"{config.model}: GSM8K accuracy={accuracy:.3f}, "
f"acceptance_len={acceptance_len:.2f}"
)
assert accuracy >= config.expected_accuracy
assert acceptance_len >= config.expected_acceptance_len
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,129 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline
from vllm import LLM
from vllm.distributed import cleanup_dist_env_and_memory
from ...utils import compute_acceptance_len, compute_acceptance_rate
def test_gemma4_dspark_correctness_and_acceptance_rate(
monkeypatch: pytest.MonkeyPatch,
):
"""
E2E test for Gemma4 DSpark speculative decoding: acceptance rate/length
regression coverage plus GSM8K correctness, at temperature=1.0 to exercise
the probabilistic draft-sampling/rejection-sampling path (not just greedy).
gemma-4-12B is instruct-tuned, so GSM8K is run through the chat template
(use_chat_completions=True; raw few-shot completion collapses to a few
percent). Reference: measured over 5 runs of 200 GSM8K questions at
temperature=1.0 (prefix caching disabled):
accuracy: min=0.900 max=0.955 mean=0.937
acceptance_rate: min=0.578 max=0.595 mean=0.588
acceptance_len: min=5.044 max=5.167 mean=5.116
Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
"""
monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
spec_llm = LLM(
model="RedHatAI/gemma-4-12B-it-NVFP4",
trust_remote_code=True,
speculative_config={
"method": "dspark",
"model": "deepseek-ai/dspark_gemma4_12b_block7",
"num_speculative_tokens": 7,
"draft_sample_method": "probabilistic",
},
max_model_len=4096,
max_num_seqs=32,
gpu_memory_utilization=0.85,
enforce_eager=True,
enable_prefix_caching=False,
disable_log_stats=False,
)
try:
results = evaluate_gsm8k_offline(
spec_llm, num_questions=200, temperature=1.0, use_chat_completions=True
)
gsm8k_accuracy = results["accuracy"]
metrics = spec_llm.get_metrics()
acceptance_rate = compute_acceptance_rate(metrics)
acceptance_len = compute_acceptance_len(metrics)
print(
f"Gemma4 DSpark acceptance_rate={acceptance_rate:.2f}, "
f"acceptance_len={acceptance_len:.2f}, "
f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
)
assert acceptance_rate >= 0.58 * 0.9
assert acceptance_len >= 5.0 * 0.9
assert gsm8k_accuracy >= 0.92 * 0.9
finally:
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
@pytest.fixture
def dspark_config():
target_model = "Qwen/Qwen3-4B-FP8"
draft_model = "deepseek-ai/dspark_qwen3_4b_block7"
return dict(
model=target_model,
trust_remote_code=True,
speculative_config={
"method": "dspark",
"model": draft_model,
"num_speculative_tokens": 7,
"attention_backend": "FLASH_ATTN",
"draft_sample_method": "probabilistic",
},
max_model_len=4096,
disable_log_stats=False,
)
def test_dspark_correctness_and_acceptance_rate(dspark_config):
"""
E2E test for DSpark speculative decoding: acceptance rate/length
regression coverage plus GSM8K correctness, at temperature=1.0 to
exercise the probabilistic draft-sampling/rejection-sampling path
(not just greedy).
Uses Qwen/Qwen3-4B-FP8 as target with the dspark_qwen3_4b_block7 draft
model. Reference: measured over 12 runs of the full GSM8K set at
temperature=1.0 (prefix caching disabled to avoid cross-run reuse):
accuracy: min=0.782 max=0.814 mean=0.801
acceptance_rate: min=0.418 max=0.434 mean=0.428
acceptance_len: min=3.928 max=4.037 mean=3.994
Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
"""
spec_llm = LLM(**dspark_config)
results = evaluate_gsm8k_offline(spec_llm, temperature=1.0)
gsm8k_accuracy = results["accuracy"]
metrics = spec_llm.get_metrics()
acceptance_rate = compute_acceptance_rate(metrics)
acceptance_len = compute_acceptance_len(metrics)
print(
f"DSpark acceptance_rate={acceptance_rate:.2f}, "
f"acceptance_len={acceptance_len:.2f}, "
f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
)
assert acceptance_rate >= 0.428 * 0.9
assert acceptance_len >= 3.994 * 0.9
assert gsm8k_accuracy >= 0.801 * 0.9
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,50 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from ...utils import compute_acceptance_rate
def test_medusa_acceptance_rate(
sampling_config: SamplingParams,
):
"""Verify a trained Medusa checkpoint achieves nonzero acceptance rate.
Uses the canonical FasterDecoding vicuna-7b checkpoint to confirm the
speculation path actually accepts tokens — unlike test_medusa_correctness,
which uses a random head and only validates output correctness.
"""
target_model = "lmsys/vicuna-7b-v1.3"
medusa_model = "FasterDecoding/medusa-vicuna-7b-v1.3"
prompts = _build_gsm8k_prompts(num_questions=10, num_shots=1)[0]
spec_llm = LLM(
model=target_model,
speculative_config={
"method": "medusa",
"model": medusa_model,
"num_speculative_tokens": 3,
},
max_model_len=1024,
enforce_eager=True,
disable_log_stats=False,
)
spec_llm.generate(prompts, sampling_config)
metrics = spec_llm.get_metrics()
acceptance_rate = compute_acceptance_rate(metrics)
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
min_acceptance_rate = 0.198
print(f"Medusa acceptance rate: {acceptance_rate:.4f} (min {min_acceptance_rate})")
# Regression guard at 90% of the measured baseline.
assert acceptance_rate >= min_acceptance_rate, (
f"Medusa acceptance rate {acceptance_rate:.4f} below min {min_acceptance_rate}"
)
@@ -0,0 +1,36 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from ..utils import run_acceptance_length_eval
@pytest.mark.parametrize("use_mrv2", [False, True])
def test_gemma4_mtp_acceptance_lengths(
monkeypatch: pytest.MonkeyPatch,
use_mrv2: bool,
):
run_acceptance_length_eval(
monkeypatch,
spec_config={
"model": "google/gemma-4-E4B-it",
"trust_remote_code": True,
"speculative_config": {
"method": "mtp",
"model": "google/gemma-4-E4B-it-assistant",
"num_speculative_tokens": 2,
"max_model_len": 32768,
},
"max_model_len": 32768,
"limit_mm_per_prompt": {"image": 0, "audio": 0},
"disable_log_stats": False,
},
expected_acceptance_lengths={
"mt-bench": 2.28,
"humaneval": 2.68,
"gsm8k": 2.67 * 0.975,
},
chat_template_kwargs={},
use_mrv2=use_mrv2,
)
@@ -0,0 +1,58 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from tests.utils import single_gpu_only
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from ...utils import compute_acceptance_len, get_test_prompts
@single_gpu_only
def test_synthetic_acceptance_rate():
"""Verify that synthetic rejection sampling produces an acceptance
length close to the requested mean acceptance length."""
num_spec_tokens = 3
expected_acceptance_len = 1.875
tolerance = 0.15
spec_llm = LLM(
model="meta-llama/Llama-3.2-1B-Instruct",
trust_remote_code=True,
speculative_config={
"method": "eagle3",
"model": "nm-testing/Llama3_2_1B_speculator.eagle3",
"num_speculative_tokens": num_spec_tokens,
"max_model_len": 2048,
"rejection_sample_method": "synthetic",
"synthetic_acceptance_length": expected_acceptance_len,
},
max_model_len=2048,
enforce_eager=True,
disable_log_stats=False,
)
test_prompts = get_test_prompts(mm_enabled=False, num_prompts=50)
spec_llm.chat(
test_prompts,
SamplingParams(temperature=0, max_tokens=64, ignore_eos=True),
)
metrics = spec_llm.get_metrics()
acceptance_len = compute_acceptance_len(metrics)
print(
f"Synthetic acceptance length: {acceptance_len:.3f}"
f" (expected={expected_acceptance_len:.3f},"
f" tolerance=±{tolerance})"
)
assert abs(acceptance_len - expected_acceptance_len) <= tolerance, (
f"Synthetic acceptance length {acceptance_len:.3f} is not within"
f" ±{tolerance} of expected {expected_acceptance_len:.3f}"
)
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
@@ -0,0 +1,105 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Any
import pytest
import torch
from tqdm import tqdm
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from ..utils import compute_acceptance_len
def load_and_process_dataset(data_name: str):
from datasets import load_dataset
if data_name == "gsm8k":
dataset = load_dataset("openai/gsm8k", "main", split="test")
prompt_fmt = (
"{question}\nPlease reason step by step,"
" and put your final answer within \\boxed{{}}."
)
dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
elif data_name == "mt-bench":
dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train")
dataset = dataset.map(lambda x: {"turns": x["prompt"]})
elif data_name == "humaneval":
dataset = load_dataset("openai/openai_humaneval", split="test")
prompt_fmt = (
"Write a solution to the following problem and make sure"
" that it passes the tests:\n```python\n{prompt}\n```"
)
dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
return dataset
def run_acceptance_length_eval(
monkeypatch: pytest.MonkeyPatch,
spec_config: dict[str, Any],
expected_acceptance_lengths: dict[str, float],
chat_template_kwargs: dict[str, Any],
use_mrv2: bool,
):
"""
E2E acceptance-rate validation for speculative decoding.
Drives one or more datasets (keyed in ``expected_acceptance_lengths``)
through the spec decode engine and asserts the mean acceptance length
stays within tolerance of the reference figure for each dataset.
"""
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
spec_llm = LLM(**spec_config)
max_prompts_per_dataset = 200 # mt-bench has 80, humaneval has 164, truncates gsm8k
tokenizer = spec_llm.get_tokenizer()
for dataset_name, expected_len in expected_acceptance_lengths.items():
dataset = load_and_process_dataset(dataset_name)
prev_metrics = None
acceptance_lengths = []
for i in tqdm(
range(min(max_prompts_per_dataset, len(dataset))),
desc=f"Processing {dataset_name}",
):
user_content = dataset[i]["turns"][0]
prompt_text = tokenizer.apply_chat_template(
[{"role": "user", "content": user_content}],
tokenize=False,
add_generation_prompt=True,
**chat_template_kwargs,
)
# Greedy (temp=0) so acceptance length is deterministic and comparable
# across runs.
spec_llm.generate(
[prompt_text],
SamplingParams(temperature=0, max_tokens=2048),
use_tqdm=False,
)
current_metrics = spec_llm.get_metrics()
acceptance_len = compute_acceptance_len(current_metrics, prev_metrics)
prev_metrics = current_metrics
acceptance_lengths.append(acceptance_len)
mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths)
# Fairly tight tolerance of 95% against the reference figures,
# watching for regressions. Can be relaxed if test is flaky but be sure to
# check for genuine issues such as #40727.
expected_len = expected_len * 0.95
print(
f"acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}"
f" (expected at least {expected_len:.2f})"
)
assert mean_acceptance_length >= expected_len, (
f"acceptance_len for {dataset_name} is below expected threshold: "
f"{mean_acceptance_length:.2f} < {expected_len:.2f}"
)
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
+25
View File
@@ -0,0 +1,25 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm import SamplingParams
from .utils import greedy_sampling
@pytest.fixture
def sampling_config() -> SamplingParams:
return greedy_sampling()
@pytest.fixture
def model_name() -> str:
return "meta-llama/Llama-3.1-8B-Instruct"
@pytest.fixture(autouse=True)
def reset_torch_dynamo():
yield
torch._dynamo.reset()
@@ -0,0 +1,388 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
import pytest
import torch
from tests.utils import multi_gpu_only, single_gpu_only
from vllm import LLM, SamplingParams
from vllm.config import VllmConfig, replace
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.engine.arg_utils import EngineArgs
from ..utils import (
Messages,
_skip_if_insufficient_gpus_for_tp,
compute_acceptance_len,
compute_acceptance_rate,
evaluate_llm_for_gsm8k,
get_instruct_coder_messages,
get_test_prompts,
greedy_sampling,
stochastic_sampling,
)
class AsyncSchedulingNotEnabledError(AssertionError):
"""Raised when draft-model spec decode does not enable async scheduling."""
@dataclass
class ArgsTest:
target_model: str
draft_model: str
sampling_config: SamplingParams
num_speculative_tokens: int
expected_acceptance_rate: float
expected_acceptance_len: float
expected_gsm8k_accuracy: float = 0.0 # skip by default
# Defaults
enforce_eager: bool = True
parallel_drafting: bool = False
target_tensor_parallel_size: int = 1
draft_tensor_parallel_size: int = 1
max_model_len: int = 2048
gpu_memory_utilization: float = 0.5
dataset: str = "test_prompts"
num_prompts: int = 100
def get_messages(dataset: str, n: int) -> list[Messages]:
if dataset == "test_prompts":
return get_test_prompts(mm_enabled=False, num_prompts=n)
if dataset == "likaixin/InstructCoder":
return get_instruct_coder_messages(n=n)
raise NotImplementedError(f"Dataset '{dataset}' not implemented")
def some_high_acceptance_metrics() -> dict:
return {
"sampling_config": greedy_sampling(),
"num_speculative_tokens": 3,
"expected_acceptance_len": 3.4, # ref: 3.75
"expected_acceptance_rate": 0.8, # ref: 0.9
}
cases = [
# Same model for draft and target, greedy sampling.
ArgsTest(
target_model="Qwen/Qwen3-0.6B",
draft_model="Qwen/Qwen3-0.6B",
sampling_config=greedy_sampling(),
num_speculative_tokens=3, # K
expected_acceptance_len=0.98 * (3 + 1), # epsilon discount of K + 1
expected_acceptance_rate=0.98, # slight epsilon
expected_gsm8k_accuracy=0.25, # ref: 35-40%
),
# Smaller draft model, stochastic sampling.
ArgsTest(
target_model="Qwen/Qwen3-1.7B",
draft_model="Qwen/Qwen3-0.6B",
sampling_config=stochastic_sampling(),
num_speculative_tokens=3,
expected_acceptance_len=3.4, # ref: 3.7
expected_acceptance_rate=0.80, # ref: 0.90
expected_gsm8k_accuracy=0.5, # ref: 60%. Note gsm8k always runs greedy sampling
),
]
@pytest.mark.parametrize("args", cases)
@pytest.mark.parametrize("enforce_eager", [True, False])
@single_gpu_only
# TODO: Fix async_scheduling & engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
@pytest.mark.xfail(
raises=AsyncSchedulingNotEnabledError,
reason="draft_model does not yet enable async_scheduling: issue #38929",
)
def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
args.enforce_eager = enforce_eager
assert_draft_model_correctness(args)
@single_gpu_only
# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
@pytest.mark.xfail(
raises=AsyncSchedulingNotEnabledError,
reason="draft_model does not yet enable async_scheduling: issue #38929",
)
def test_draft_model_realistic_example():
args = ArgsTest(
target_model="Qwen/Qwen3-1.7B",
draft_model="Qwen/Qwen3-0.6B",
dataset="likaixin/InstructCoder",
num_speculative_tokens=3,
sampling_config=greedy_sampling(),
enforce_eager=False,
expected_acceptance_len=2.6, # ref: 2.86
expected_acceptance_rate=0.5, # ref: 0.62
)
assert_draft_model_correctness(args)
@single_gpu_only
# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
@pytest.mark.xfail(
raises=AsyncSchedulingNotEnabledError,
reason="draft_model does not yet enable async_scheduling: issue #38929",
)
def test_draft_model_parallel_drafting():
args = ArgsTest(
target_model="Qwen/Qwen3-1.7B",
draft_model="amd/PARD-Qwen3-0.6B",
dataset="likaixin/InstructCoder",
num_speculative_tokens=3,
sampling_config=greedy_sampling(),
parallel_drafting=True,
enforce_eager=False,
expected_acceptance_len=2.3, # ref: 2.52
expected_acceptance_rate=0.4, # ref: 0.51
)
assert_draft_model_correctness(args)
@pytest.mark.parametrize(
"models",
[
# target_model, draft_model
("Qwen/Qwen3-1.7B-FP8", "Qwen/Qwen3-0.6B"), # target quantized
("Qwen/Qwen3-1.7B", "Qwen/Qwen3-0.6B-FP8"), # draft quantized
],
ids=["target_quantized", "draft_quantized"],
)
@pytest.mark.parametrize("enforce_eager", [True, False])
@single_gpu_only
# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
@pytest.mark.xfail(
raises=AsyncSchedulingNotEnabledError,
reason="draft_model does not yet enable async_scheduling: issue #38929",
)
def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
tgt_model, draft_model = models
sd_case = ArgsTest(
target_model=tgt_model,
draft_model=draft_model,
**some_high_acceptance_metrics(),
enforce_eager=enforce_eager,
)
assert_draft_model_correctness(sd_case)
@multi_gpu_only(num_gpus=2)
# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
@pytest.mark.xfail(
raises=AsyncSchedulingNotEnabledError,
reason="draft_model does not yet enable async_scheduling: issue #38929",
)
def test_draft_model_tensor_parallelism():
"""Ensure spec decode works when running with TP > 1."""
_skip_if_insufficient_gpus_for_tp(2)
sd_case = ArgsTest(
target_model="Qwen/Qwen3-1.7B",
target_tensor_parallel_size=2,
draft_model="Qwen/Qwen3-0.6B",
draft_tensor_parallel_size=2,
**some_high_acceptance_metrics(),
enforce_eager=False,
expected_gsm8k_accuracy=0.5,
)
assert_draft_model_correctness(sd_case)
@multi_gpu_only(num_gpus=2)
def test_draft_model_engine_args_tensor_parallelism():
"""Ensure the vllm_config for the draft model is created correctly,
and independently of the target model (quantization, TP, etc.)"""
_skip_if_insufficient_gpus_for_tp(2)
engine_args = EngineArgs(
model="Qwen/Qwen3-1.7B-FP8", # <<< tgt quantized
tensor_parallel_size=2,
speculative_config={
"model": "Qwen/Qwen3-0.6B", # <<< draft not quantized
"method": "draft_model",
"num_speculative_tokens": 3,
"draft_tensor_parallel_size": 1, # <<< valid arg name
},
)
target_config: VllmConfig = engine_args.create_engine_config()
assert target_config.parallel_config.tensor_parallel_size == 2
assert target_config.quant_config.get_name() == "fp8"
speculative_config = target_config.speculative_config
draft_config: VllmConfig = replace(
target_config,
quant_config=None,
parallel_config=replace(
speculative_config.draft_parallel_config,
rank=target_config.parallel_config.rank,
),
model_config=speculative_config.draft_model_config,
)
assert draft_config.parallel_config.tensor_parallel_size == 1
assert draft_config.quant_config is None
def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig:
"""Replicate SpecDecodeBaseProposer._create_draft_vllm_config logic
so we can test it without instantiating a full proposer."""
spec_cfg = vllm_config.speculative_config
if spec_cfg.moe_backend is not None:
return replace(
vllm_config,
kernel_config=replace(
vllm_config.kernel_config,
moe_backend=spec_cfg.moe_backend,
),
)
return vllm_config
def test_draft_model_moe_backend_override():
"""When moe_backend is set in speculative_config, the draft VllmConfig
should use it while the target keeps its own setting."""
engine_args = EngineArgs(
model="Qwen/Qwen3-1.7B",
tensor_parallel_size=1,
moe_backend="flashinfer_trtllm",
speculative_config={
"model": "Qwen/Qwen3-0.6B",
"method": "draft_model",
"num_speculative_tokens": 3,
"moe_backend": "triton",
},
)
tgt_config: VllmConfig = engine_args.create_engine_config()
assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm"
assert tgt_config.speculative_config.moe_backend == "triton"
draft_config = _apply_draft_moe_backend(tgt_config)
assert draft_config.kernel_config.moe_backend == "triton"
# Target config must be unaffected.
assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm"
def test_draft_model_moe_backend_inherits_target():
"""When moe_backend is not set in speculative_config, the draft should
inherit the target's moe_backend."""
engine_args = EngineArgs(
model="Qwen/Qwen3-1.7B",
tensor_parallel_size=1,
moe_backend="flashinfer_cutlass",
speculative_config={
"model": "Qwen/Qwen3-0.6B",
"method": "draft_model",
"num_speculative_tokens": 3,
},
)
tgt_config: VllmConfig = engine_args.create_engine_config()
assert tgt_config.kernel_config.moe_backend == "flashinfer_cutlass"
assert tgt_config.speculative_config.moe_backend is None
draft_config = _apply_draft_moe_backend(tgt_config)
assert draft_config.kernel_config.moe_backend == "flashinfer_cutlass"
assert draft_config is tgt_config
def test_draft_model_moe_backend_default_auto():
"""When neither target nor draft set moe_backend explicitly, both should
default to 'auto'."""
engine_args = EngineArgs(
model="Qwen/Qwen3-1.7B",
tensor_parallel_size=1,
speculative_config={
"model": "Qwen/Qwen3-0.6B",
"method": "draft_model",
"num_speculative_tokens": 3,
},
)
tgt_config: VllmConfig = engine_args.create_engine_config()
assert tgt_config.kernel_config.moe_backend == "auto"
assert tgt_config.speculative_config.moe_backend is None
draft_config = _apply_draft_moe_backend(tgt_config)
assert draft_config.kernel_config.moe_backend == "auto"
assert draft_config is tgt_config
def test_draft_model_engine_args_rejects_invalid_tp_argname():
"""The user should pass "draft_tensor_parallel_size" rather than
"tensor_parallel_size". We enforce this with validation."""
engine_args = EngineArgs(
model="Qwen/Qwen3-1.7B",
tensor_parallel_size=1,
speculative_config={
"model": "Qwen/Qwen3-0.6B",
"method": "draft_model",
"num_speculative_tokens": 3,
"tensor_parallel_size": 1, # <<< invalid arg name
},
)
with pytest.raises(ValueError):
engine_args.create_engine_config()
def assert_draft_model_correctness(args: ArgsTest):
"""Compare the outputs using and not using speculative decoding.
In the greedy decoding case, the outputs must match EXACTLY."""
test_prompts: list[Messages] = get_messages(
dataset=args.dataset, n=args.num_prompts
)
spec_llm = LLM(
model=args.target_model,
speculative_config={
"model": args.draft_model,
"method": "draft_model",
"num_speculative_tokens": args.num_speculative_tokens,
"max_model_len": args.max_model_len,
"enforce_eager": args.enforce_eager,
"draft_tensor_parallel_size": args.draft_tensor_parallel_size,
"parallel_drafting": args.parallel_drafting,
},
max_num_seqs=100, # limit cudagraph capture runtime
max_model_len=args.max_model_len,
gpu_memory_utilization=args.gpu_memory_utilization,
tensor_parallel_size=args.target_tensor_parallel_size,
enforce_eager=args.enforce_eager,
disable_log_stats=False, # enables get_metrics()
)
# we don't check the outputs, only check the metrics
spec_llm.chat(test_prompts, args.sampling_config)
metrics = spec_llm.get_metrics()
acceptance_rate: float = compute_acceptance_rate(metrics)
acceptance_len: float = compute_acceptance_len(metrics)
# Need to evaluate after getting metrics to avoid polluting the AR
evaluate_llm_for_gsm8k(
spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy
)
print(
f"spec-decode: target={args.target_model}, draft={args.draft_model}, "
f"temperature={args.sampling_config.temperature:.2f}, "
f"acceptance_rate={acceptance_rate:.2f}, "
f"acceptance_len={acceptance_len:.2f}, "
)
assert acceptance_rate >= args.expected_acceptance_rate
assert acceptance_len >= args.expected_acceptance_len
# draft_model supports async scheduling; assert it is active by default.
# Raise AsyncSchedulingNotEnabledError (a subclass of AssertionError) so that
# @pytest.mark.xfail(raises=AsyncSchedulingNotEnabledError) catches only this
# specific failure — leaving all other assertion failures (e.g. correctness or
# acceptance-rate checks above) visible as real test failures.
has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
del spec_llm # CLEANUP
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
if not has_async:
raise AsyncSchedulingNotEnabledError(
"Expected async_scheduling=True for draft_model spec decode, got False."
" See https://github.com/vllm-project/vllm/issues/38929"
)
@@ -0,0 +1,167 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from tests.utils import (
get_attn_backend_list_based_on_platform,
single_gpu_only,
)
from vllm import SamplingParams
from vllm.platforms import current_platform
from .utils import _run_eagle_correctness
@single_gpu_only
@pytest.mark.skipif(
current_platform.is_device_capability_family(100),
reason="DeepSeek head_dim=192 not supported on SM100/SM110 (Blackwell)",
)
@pytest.mark.parametrize(
[
"model_setup",
"mm_enabled",
"enable_chunked_prefill",
"model_impl",
"expected_accuracy_threshold",
],
[
(
(
"eagle",
"eagle618/deepseek-v3-random",
"eagle618/eagle-deepseek-v3-random",
1,
),
False,
False,
"auto",
0.0,
),
],
ids=["deepseek_eagle"],
)
@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
def test_eagle_correctness_light(
monkeypatch: pytest.MonkeyPatch,
sampling_config: SamplingParams,
model_setup: tuple[str, str, str, int],
mm_enabled: bool,
expected_accuracy_threshold: float,
enable_chunked_prefill: bool,
model_impl: str,
attn_backend: str,
):
_run_eagle_correctness(
monkeypatch,
sampling_config,
model_setup,
mm_enabled,
expected_accuracy_threshold,
enable_chunked_prefill,
model_impl,
attn_backend,
)
@single_gpu_only
@pytest.mark.parametrize(
[
"model_setup",
"mm_enabled",
"enable_chunked_prefill",
"model_impl",
"expected_accuracy_threshold",
],
[
(
("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
False,
False,
"auto",
0.8,
),
pytest.param(
("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
False,
False,
"transformers",
0.8,
# TODO(hmellor): figure out why memory usage is so high
marks=pytest.mark.skip(
reason="Feature is experimental and uses too much memory in CI",
),
),
pytest.param(
(
"eagle3",
"Qwen/Qwen3-VL-8B-Instruct",
"taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
1,
),
False,
False,
"auto",
0.8,
marks=pytest.mark.skip(
reason="architecture of its eagle3 is LlamaForCausalLMEagle3"
),
),
pytest.param(
(
"eagle3",
"Qwen/Qwen2.5-VL-7B-Instruct",
"Rayzl/qwen2.5-vl-7b-eagle3-sgl",
1,
),
False,
False,
"auto",
0.7,
marks=pytest.mark.skip(
reason="Skipping due to its head_dim not being a multiple of 32"
),
),
(
(
"eagle3",
"meta-llama/Llama-3.1-8B-Instruct",
"yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
1,
),
False,
False,
"auto",
0.7,
),
],
ids=[
"qwen3_eagle3",
"qwen3_eagle3-transformers",
"qwen3_vl_eagle3",
"qwen2_5_vl_eagle3",
"llama3_eagle3",
],
)
@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
def test_eagle_correctness_medium(
monkeypatch: pytest.MonkeyPatch,
sampling_config: SamplingParams,
model_setup: tuple[str, str, str, int],
mm_enabled: bool,
expected_accuracy_threshold: float,
enable_chunked_prefill: bool,
model_impl: str,
attn_backend: str,
):
_run_eagle_correctness(
monkeypatch,
sampling_config,
model_setup,
mm_enabled,
expected_accuracy_threshold,
enable_chunked_prefill,
model_impl,
attn_backend,
)
+137
View File
@@ -0,0 +1,137 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from tests.utils import wait_for_rocm_memory_to_settle
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.platforms import current_platform
from ..utils import (
_skip_if_insufficient_gpus_for_tp,
evaluate_llm_for_gsm8k,
get_test_prompts,
)
def _run_eagle_correctness(
monkeypatch: pytest.MonkeyPatch,
sampling_config: SamplingParams,
model_setup: tuple[str, str, str, int],
mm_enabled: bool,
expected_accuracy_threshold: float,
enable_chunked_prefill: bool,
model_impl: str,
attn_backend: str,
):
"""
Compare the outputs of an original LLM and a speculative LLM
which should be the same when using eagle speculative decoding.
"""
if model_impl == "transformers":
import transformers
from packaging.version import Version
installed = Version(transformers.__version__)
required = Version("5.0.0")
if installed < required:
pytest.skip(
"Eagle3 with the Transformers modeling backend requires "
f"transformers>={required}, but got {installed}"
)
test_prompts = get_test_prompts(mm_enabled)
if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN":
if current_platform.is_rocm():
print(
"FLASH_ATTN for spec_decode not supported on "
"ROCm currently. Changing to FLEX_ATTENTION backend."
)
attention_config = {"backend": "FLEX_ATTENTION"}
else:
attention_config = None
else:
attention_config = {"backend": attn_backend}
if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm():
pytest.skip(
"TRITON_ATTN does not support "
"multi-token eagle spec decode on current platform"
)
with monkeypatch.context() as m:
m.setenv("VLLM_MLA_DISABLE", "1")
if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
if "deepseek" in model_setup[1].lower():
m.setenv("VLLM_ROCM_USE_AITER", "1")
m.delenv("VLLM_MLA_DISABLE", raising=False)
attention_config = {"backend": "ROCM_AITER_MLA"}
else:
m.setenv("VLLM_ROCM_USE_AITER", "1")
method, model_name, spec_model_name, tp_size = model_setup
_skip_if_insufficient_gpus_for_tp(tp_size)
max_model_len = 2048
max_num_batched_tokens = 128 if enable_chunked_prefill else max_model_len
ref_llm = LLM(
model=model_name,
max_model_len=max_model_len,
tensor_parallel_size=tp_size,
attention_config=attention_config,
)
evaluate_llm_for_gsm8k(
ref_llm, expected_accuracy_threshold=expected_accuracy_threshold
)
ref_outputs = ref_llm.chat(test_prompts, sampling_config)
del ref_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
# ROCm frees VRAM lazily; wait so the spec engine started right after
# does not OOM on its startup memory guard.
wait_for_rocm_memory_to_settle()
spec_llm = LLM(
model=model_name,
trust_remote_code=True,
tensor_parallel_size=tp_size,
speculative_config={
"method": method,
"model": spec_model_name,
"num_speculative_tokens": 3,
"max_model_len": max_model_len,
},
max_model_len=max_model_len,
max_num_batched_tokens=max_num_batched_tokens,
enable_chunked_prefill=enable_chunked_prefill,
model_impl=model_impl,
attention_config=attention_config,
)
# EAGLE/EAGLE3 supports async scheduling; assert it is active by default.
assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
evaluate_llm_for_gsm8k(
spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
)
spec_outputs = spec_llm.chat(test_prompts, sampling_config)
matches = 0
misses = 0
for ref_output, spec_output in zip(ref_outputs, spec_outputs):
if ref_output.outputs[0].text == spec_output.outputs[0].text:
matches += 1
else:
misses += 1
print(f"ref_output: {ref_output.outputs[0].text}")
print(f"spec_output: {spec_output.outputs[0].text}")
assert matches > int(0.6 * len(ref_outputs))
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
# ROCm frees VRAM lazily; wait so the next parametrization's engine does
# not OOM on its startup memory guard.
wait_for_rocm_memory_to_settle()
+155
View File
@@ -0,0 +1,155 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
from typing import Any
import pytest
import torch
from tests.utils import single_gpu_only
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.platforms import current_platform
from ..utils import (
_skip_if_insufficient_gpus_for_tp,
evaluate_llm_for_gsm8k,
get_test_prompts,
)
@pytest.mark.parametrize(
["model_setup", "mm_enabled", "expected_accuracy_threshold"],
[
(("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False, 0.5), # ref: 65%-70%
pytest.param(
("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1),
False,
0.0,
marks=pytest.mark.skipif(
current_platform.is_device_capability_family(100),
reason="DeepSeek MTP: TRTLLM MoE top_k check fails on Blackwell",
),
), # dummy model
(
("mtp", "Qwen/Qwen3.5-0.8B-Base", 1),
False,
0.20,
), # hybrid + MTP, ref: ~34%-35%
(
("mtp", "google/gemma-4-E4B-it", 1, "google/gemma-4-E4B-it-assistant"),
False,
0.50,
), # gemma4 MTP with assistant model, ref: ~62%
],
ids=["mimo", "deepseek", "qwen3_5-hybrid", "gemma4-e4b"],
)
@single_gpu_only
def test_mtp_correctness(
monkeypatch: pytest.MonkeyPatch,
sampling_config: SamplingParams,
model_setup: tuple[str, str, int] | tuple[str, str, int, str],
mm_enabled: bool,
expected_accuracy_threshold: float,
):
"""
Compare the outputs of a original LLM and a speculative LLM
which should be the same when using MTP speculative decoding. Due to some variance
in the engine, it is possible for some outputs to differ, so we expect that at least
6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed
reference threshold for each model.
"""
# Generate test prompts inside the function instead of using fixture
test_prompts = get_test_prompts(mm_enabled)
with monkeypatch.context() as m:
m.setenv("VLLM_MLA_DISABLE", "1")
if len(model_setup) == 4:
method, model_name, tp_size, draft_model = model_setup
else:
method, model_name, tp_size = model_setup
draft_model = None
_skip_if_insufficient_gpus_for_tp(tp_size)
if "Qwen3.5" in model_name and os.environ.get("VLLM_USE_V2_MODEL_RUNNER"):
pytest.skip(
"Model Runner V2 does not yet support hybrid models "
"(Qwen3.5 mixes Mamba-style GDN with attention layers)."
)
attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
# Skip multimodal profiling for models that don't need it in this test.
extra_kwargs: dict[str, Any] = {}
if "Qwen3.5" in model_name:
extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0}
elif "gemma-4" in model_name:
extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "audio": 0}
if draft_model is not None and "gemma-4" in draft_model:
import transformers
from packaging.version import Version
if Version(transformers.__version__) < Version("5.8.0"):
pytest.skip(
"Gemma4 MTP assistant requires transformers>=5.8.0, "
f"got {transformers.__version__}"
)
ref_llm = LLM(
model=model_name,
max_model_len=2048,
tensor_parallel_size=tp_size,
trust_remote_code=True,
attention_backend=attn_backend,
**extra_kwargs,
)
ref_outputs = ref_llm.chat(test_prompts, sampling_config)
evaluate_llm_for_gsm8k(
ref_llm, expected_accuracy_threshold=expected_accuracy_threshold
)
del ref_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
speculative_config: dict[str, Any] = {
"method": method,
"num_speculative_tokens": 1,
"max_model_len": 2048,
}
if draft_model is not None:
speculative_config["model"] = draft_model
speculative_config["num_speculative_tokens"] = 2
spec_llm = LLM(
model=model_name,
trust_remote_code=True,
tensor_parallel_size=tp_size,
speculative_config=speculative_config,
max_model_len=2048,
attention_backend=attn_backend,
**extra_kwargs,
)
# MTP supports async scheduling; assert it is active by default.
assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
evaluate_llm_for_gsm8k(
spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
)
spec_outputs = spec_llm.chat(test_prompts, sampling_config)
matches = 0
misses = 0
for ref_output, spec_output in zip(ref_outputs, spec_outputs):
if ref_output.outputs[0].text == spec_output.outputs[0].text:
matches += 1
else:
misses += 1
print(f"ref_output: {ref_output.outputs[0].text}")
print(f"spec_output: {spec_output.outputs[0].text}")
# Heuristic: expect at least 80% of the prompts to match exactly
# Upon failure, inspect the outputs to check for inaccuracy.
assert matches > int(0.8 * len(ref_outputs))
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
@@ -0,0 +1,147 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from tests.utils import single_gpu_only
from vllm import LLM, SamplingParams
from vllm.config import CompilationConfig
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.platforms import current_platform
from ..utils import evaluate_llm_for_gsm8k, get_test_prompts
@pytest.fixture
def disable_vllm_compile_cache_on_rocm(request: pytest.FixtureRequest) -> None:
if current_platform.is_rocm():
request.getfixturevalue("disable_vllm_compile_cache")
@pytest.mark.parametrize(
"speculative_config",
[
{
"method": "ngram",
"prompt_lookup_max": 5,
"prompt_lookup_min": 3,
"num_speculative_tokens": 3,
},
{
"method": "suffix",
"suffix_decoding_max_spec_factor": 2.0,
},
],
)
@pytest.mark.usefixtures("disable_vllm_compile_cache_on_rocm")
@single_gpu_only
def test_ngram_and_suffix_correctness(
speculative_config: dict,
model_name: str,
vllm_runner,
):
with vllm_runner(
model_name,
# Keep LLM defaults; VllmRunner only provides lifecycle cleanup here.
trust_remote_code=False,
enable_chunked_prefill=None,
speculative_config=speculative_config,
max_model_len=4096,
# Preserve LLM's default compilation/cudagraph configuration. Without
# this, VllmRunner injects its reduced test-only capture sizes.
compilation_config=CompilationConfig(),
) as runner:
evaluate_llm_for_gsm8k(runner.llm)
@pytest.mark.parametrize("async_scheduling", [True], ids=["async"])
@single_gpu_only
def test_ngram_gpu_default_with_async_scheduling(
async_scheduling: bool,
):
"""
Test ngram_gpu speculative decoding (k=3) correctness with and without
async scheduling, validated via GSM8K accuracy.
Uses Qwen/Qwen3-8B (ref GSM8K accuracy: 87%-92%).
"""
qwen3_model = "Qwen/Qwen3-8B"
spec_llm = LLM(
model=qwen3_model,
speculative_config={
"method": "ngram_gpu",
"prompt_lookup_max": 3,
"prompt_lookup_min": 2,
"num_speculative_tokens": 2,
},
max_model_len=4096,
async_scheduling=async_scheduling,
)
# Assert the resolved async_scheduling config matches what was requested.
assert (
spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
== async_scheduling
)
evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8)
del spec_llm
cleanup_dist_env_and_memory()
@single_gpu_only
def test_suffix_decoding_acceptance(
monkeypatch: pytest.MonkeyPatch,
sampling_config: SamplingParams,
model_name: str,
):
"""
Check that suffix decoding caching takes effect and improves acceptance
lengths and acceptance rates over multiple runs of the same prompts.
"""
test_prompts = get_test_prompts(mm_enabled=False)
spec_llm = LLM(
model=model_name,
speculative_config={
"method": "suffix",
"suffix_decoding_max_spec_factor": 2.0,
"suffix_decoding_max_cached_requests": 1000,
},
max_model_len=1024,
disable_log_stats=False,
)
# Run several times and check that the accepted tokens increase.
num_draft = []
num_accept = []
for i in range(10): # Run multiple times to warm up the cache.
spec_llm.chat(test_prompts, sampling_config)
# Collect draft and acceptance stats.
metrics = spec_llm.get_metrics()
for metric in metrics:
if metric.name == "vllm:spec_decode_num_draft_tokens":
num_draft.append(metric.value)
if metric.name == "vllm:spec_decode_num_accepted_tokens":
num_accept.append(metric.value)
# Calculate the acceptance rates for the first and last runs.
first_accept_tokens = num_accept[0]
first_draft_tokens = num_draft[0]
first_accept_rate = first_accept_tokens / first_draft_tokens
# Take the diff since the stats are cumulative.
last_accept_tokens = num_accept[-1] - num_accept[-2]
last_draft_tokens = num_draft[-1] - num_draft[-2]
last_accept_rate = last_accept_tokens / last_draft_tokens
# Expect the acceptance length to improve.
assert first_accept_tokens < last_accept_tokens
# Expect the acceptance rate to improve.
assert first_accept_rate < last_accept_rate
# Heuristic: expect at least 80.0% acceptance rate at the end.
assert last_accept_rate > 0.80
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,104 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from tests.utils import single_gpu_only
from vllm import SamplingParams
from vllm.config import CompilationConfig
from ..utils import evaluate_llm_for_gsm8k, get_test_prompts
@pytest.mark.parametrize(
["model_path", "expected_accuracy_threshold"],
[
("RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3", 0.7), # ref: 75%-80%
("RedHatAI/Qwen3-8B-speculator.eagle3", 0.8), # ref: 87%-92%
],
ids=["llama3_eagle3_speculator", "qwen3_eagle3_speculator"],
)
@single_gpu_only
def test_speculators_model_integration(
monkeypatch: pytest.MonkeyPatch,
sampling_config: SamplingParams,
model_path: str,
expected_accuracy_threshold: float,
vllm_runner,
):
"""
Test that speculators models work with the simplified integration.
This verifies the `vllm serve <speculator-model>` use case where
speculative config is automatically detected from the model config
without requiring explicit --speculative-config argument.
Tests:
1. Speculator model is correctly detected
2. Verifier model is extracted from speculator config
3. Speculative decoding is automatically enabled
4. Text generation works correctly
5. GSM8k accuracy of the model passes a sanity check when speculative decoding on
6. Output matches reference (non-speculative) generation
"""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
# Generate test prompts
test_prompts = get_test_prompts(mm_enabled=False)
# First run: Direct speculator model (simplified integration)
with vllm_runner(
model_path,
trust_remote_code=False,
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
max_model_len=4096,
gpu_memory_utilization=0.92,
) as spec_runner:
evaluate_llm_for_gsm8k(
spec_runner.llm, expected_accuracy_threshold=expected_accuracy_threshold
)
spec_outputs = spec_runner.llm.chat(test_prompts, sampling_config)
# Verify speculative config was auto-detected
assert spec_runner.llm.llm_engine.vllm_config.speculative_config is not None, (
f"Speculative config should be auto-detected for {model_path}"
)
spec_config = spec_runner.llm.llm_engine.vllm_config.speculative_config
assert spec_config.num_speculative_tokens > 0, (
f"Expected positive speculative tokens, "
f"got {spec_config.num_speculative_tokens}"
)
# Verify draft model is set to the speculator model
assert spec_config.model == model_path, (
f"Draft model should be {model_path}, got {spec_config.model}"
)
# Extract verifier model for reference run
verifier_model = spec_runner.llm.llm_engine.vllm_config.model_config.model
# Second run: Reference without speculative decoding
with vllm_runner(
verifier_model,
trust_remote_code=False,
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
max_model_len=4096,
gpu_memory_utilization=0.92,
) as ref_runner:
ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config)
# Compare outputs
matches = sum(
1
for ref, spec in zip(ref_outputs, spec_outputs)
if ref.outputs[0].text == spec.outputs[0].text
)
# Heuristic: expect at least 66% of prompts to match exactly
assert matches >= int(0.66 * len(ref_outputs)), (
f"Only {matches}/{len(ref_outputs)} outputs matched. "
f"Expected at least {int(0.66 * len(ref_outputs))} matches."
)
@@ -1,60 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from tests.utils import large_gpu_mark
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
def _get_counter(metrics, name: str) -> float:
metric = next((m for m in metrics if m.name == name), None)
assert metric is not None, f"Missing metric: {name}"
return metric.value
@pytest.mark.slow_test
@large_gpu_mark(min_gb=80)
def test_laguna_dflash_hf_pair_smoke(monkeypatch):
"""Smoke-test the public Laguna XS-2.1 base/DFlash checkpoint pair."""
monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
llm = LLM(
model="poolside/Laguna-XS-2.1",
trust_remote_code=True,
speculative_config={
"method": "dflash",
"model": "poolside/Laguna-XS-2.1-DFlash",
"num_speculative_tokens": 15,
},
max_model_len=8192,
max_num_batched_tokens=65536,
max_num_seqs=32,
enforce_eager=True,
disable_log_stats=False,
)
try:
outputs = llm.generate(
[
"What is the capital of the United Kingdom?",
"Write a Python function that returns the square of a number.",
],
SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True),
)
assert len(outputs) == 2
assert all(output.outputs[0].text for output in outputs)
metrics = llm.get_metrics()
num_drafts = _get_counter(metrics, "vllm:spec_decode_num_drafts")
num_accepted = _get_counter(metrics, "vllm:spec_decode_num_accepted_tokens")
assert num_drafts > 0
acceptance_len = 1 + (num_accepted / num_drafts)
assert acceptance_len > 1.0
finally:
del llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
File diff suppressed because it is too large Load Diff
+158
View File
@@ -0,0 +1,158 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import random
from collections.abc import Iterable
from typing import Any
import pytest
import torch
from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline
from vllm import LLM, SamplingParams
from vllm.assets.base import VLLM_S3_BUCKET_URL
from vllm.assets.image import VLM_IMAGES_DIR
from vllm.v1.metrics.reader import Metric
def _skip_if_insufficient_gpus_for_tp(tp_size: int):
"""Skip test if available GPUs < tp_size on ROCm."""
available_gpus = torch.accelerator.device_count()
if available_gpus < tp_size:
pytest.skip(
f"Test requires {tp_size} GPUs, but only {available_gpus} available"
)
Messages = list[dict[str, Any]]
def get_test_prompts(mm_enabled: bool, num_prompts: int = 100) -> list[Messages]:
prompt_types = ["repeat", "gsm8k"]
if mm_enabled:
prompt_types.append("mm")
prompts: list[Messages] = []
num_repeat_prompts = num_prompts // len(prompt_types)
if mm_enabled:
num_gsm8k_prompts = num_prompts // len(prompt_types)
num_mm_prompts = num_prompts - num_repeat_prompts - num_gsm8k_prompts
else:
num_mm_prompts = 0
num_gsm8k_prompts = num_prompts - num_repeat_prompts
# Generate a mixed batch of prompts, some of which can be easily
# predicted by n-gram matching and some which likely cannot.
random.seed(0)
for _ in range(num_repeat_prompts):
word_choices = ["test", "temp", "hello", "where"]
word = random.choice(word_choices)
prompts.append(
[
{
"role": "user",
"content": f"""
please repeat the word '{word}' 10 times.
give no other output than the word at least ten times in a row,
in lowercase with spaces between each word and without quotes.
""",
}
]
)
prompts.extend(
[{"role": "user", "content": prompt}]
for prompt in _build_gsm8k_prompts(
num_questions=num_gsm8k_prompts, num_shots=5
)[0]
)
for _ in range(num_mm_prompts):
placeholders = [
{
"type": "image_url",
"image_url": {
"url": f"{VLLM_S3_BUCKET_URL}/{VLM_IMAGES_DIR}/stop_sign.jpg"
},
}
]
prompt = [
*placeholders,
{"type": "text", "text": "The meaning of the image is"},
]
prompts.append([{"role": "user", "content": prompt}])
return prompts
def get_instruct_coder_messages(n: int) -> list[Messages]:
from vllm.benchmarks.datasets import InstructCoderDataset
dataset = InstructCoderDataset(
dataset_path="likaixin/InstructCoder", dataset_split="train"
)
prompts: Iterable[str] = dataset.sample_prompts(n=n)
return [[{"role": "user", "content": prompt}] for prompt in prompts]
def greedy_sampling() -> SamplingParams:
return SamplingParams(temperature=0, max_tokens=10, ignore_eos=False)
def stochastic_sampling() -> SamplingParams:
return SamplingParams(temperature=1.0, max_tokens=10, ignore_eos=False)
def evaluate_llm_for_gsm8k(llm: LLM, expected_accuracy_threshold: float = 0.70) -> None:
"""Evaluate the LLM on GSM8K and check that accuracy is above a sanity threshold.
The default threshold assumes the LLM uses the same target model as the "model_name"
fixture, with max model len == 4096. Precomputed reference value is 75% to 80%
on GSM8K with greedy decoding, so we check that it's above a sanity threshold of 70%
to verify that the model is correct.
"""
if expected_accuracy_threshold <= 0.0:
print("Skipping GSM8K evaluation")
return
results = evaluate_gsm8k_offline(llm)
accuracy = results["accuracy"]
print(f"GSM8K accuracy: {accuracy:.3f}")
assert accuracy >= expected_accuracy_threshold, (
f"Expected GSM8K accuracy >= {expected_accuracy_threshold}, got {accuracy:.3f}"
)
def compute_acceptance_rate(
metrics: list[Metric], prev_metrics: list[Metric] | None = None
) -> float:
name2metric = {metric.name: metric for metric in metrics}
n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value
if n_draft_toks == 0:
return float("nan")
n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value
if prev_metrics is not None:
prev_name2metric = {metric.name: metric for metric in prev_metrics}
n_draft_toks -= prev_name2metric["vllm:spec_decode_num_draft_tokens"].value
n_accepted_toks -= prev_name2metric[
"vllm:spec_decode_num_accepted_tokens"
].value
if n_draft_toks <= 0:
return float("nan")
return n_accepted_toks / n_draft_toks
def compute_acceptance_len(
metrics: list[Metric], prev_metrics: list[Metric] | None = None
) -> float:
name2metric = {metric.name: metric for metric in metrics}
n_drafts = name2metric["vllm:spec_decode_num_drafts"].value
n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value
if n_drafts == 0:
return 1
if prev_metrics is not None:
prev_name2metric = {metric.name: metric for metric in prev_metrics}
n_drafts -= prev_name2metric["vllm:spec_decode_num_drafts"].value
n_accepted_toks -= prev_name2metric[
"vllm:spec_decode_num_accepted_tokens"
].value
if n_drafts <= 0:
return 1
return 1 + (n_accepted_toks / n_drafts)
+4 -2
View File
@@ -22,7 +22,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer
from .qwen3_dflash import DFlashQwen3Model
from .qwen3_dflash import DFlashQwen3Model, _dflash_layer_causal
from .qwen3_dspark import DSparkMarkovHead, Qwen3DSparkForCausalLM
from .utils import extract_layer_index, maybe_prefix
@@ -37,7 +37,8 @@ class Gemma4DSparkAttention(Gemma4MTPAttention):
quant_config: QuantizationConfig | None,
prefix: str,
) -> None:
is_full = config.layer_types[extract_layer_index(prefix)] == "full_attention"
layer_idx = extract_layer_index(prefix)
is_full = config.layer_types[layer_idx] == "full_attention"
head_dim = (
getattr(config, "global_head_dim", config.head_dim)
if is_full
@@ -62,6 +63,7 @@ class Gemma4DSparkAttention(Gemma4MTPAttention):
prefix=prefix,
)
self.is_kv_shared_layer = False
self.causal = _dflash_layer_causal(config, layer_idx)
self.use_k_eq_v = use_k_eq_v
self.kv_size = self.num_kv_heads * self.head_dim
attn_bias = getattr(config, "attention_bias", False)
@@ -19,6 +19,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
from vllm.compilation.backends import set_model_tag
from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal
from vllm.model_executor.models.utils import get_draft_quant_config
draft_vllm_config = replace(
vllm_config,
@@ -36,6 +37,9 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
else vllm_config.cache_config
),
)
# VllmConfig post-init restores the target's quant config because the target
# config is retained for DSpark's target-layer metadata, so we must override it.
draft_vllm_config.quant_config = get_draft_quant_config(vllm_config)
with set_model_tag("dspark_head"):
draft_model = get_model(