diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml index 89e25f23be..d0aabd6d32 100644 --- a/.buildkite/intel_jobs/engine_intel.yaml +++ b/.buildkite/intel_jobs/engine_intel.yaml @@ -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' diff --git a/.buildkite/intel_jobs/model_runner_v2_intel.yaml b/.buildkite/intel_jobs/model_runner_v2_intel.yaml index 43ab989fbb..85a23eaef0 100644 --- a/.buildkite/intel_jobs/model_runner_v2_intel.yaml +++ b/.buildkite/intel_jobs/model_runner_v2_intel.yaml @@ -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 diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index c3cfc55080..5845844dd3 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -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] diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 6099f09ab3..cb311ab6a1 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -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 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index caa56c21b3..f89abadbe0 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -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 diff --git a/.buildkite/test_areas/model_runner_v2.yaml b/.buildkite/test_areas/model_runner_v2.yaml index bc0b56f4ca..53b35ce829 100644 --- a/.buildkite/test_areas/model_runner_v2.yaml +++ b/.buildkite/test_areas/model_runner_v2.yaml @@ -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/ diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 7cde7124cd..464c4c021d 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -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/ diff --git a/tests/models/registry.py b/tests/models/registry.py index 24bc4a905a..f8d0387dff 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -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, diff --git a/tests/utils.py b/tests/utils.py index 90f12d444e..fc5ede2d99 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -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)." + ), ) diff --git a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py index 11f77492d2..3155da82ea 100644 --- a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py +++ b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py @@ -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 diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/__init__.py b/tests/v1/e2e/spec_decode/acceptance_rates/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/__init__.py b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py new file mode 100644 index 0000000000..a4c9c67a05 --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py @@ -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() diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dspark/__init__.py b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py new file mode 100644 index 0000000000..0827d80f66 --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py @@ -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() diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/__init__.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py new file mode 100644 index 0000000000..cd9c115a9d --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py @@ -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}" + ) diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py new file mode 100644 index 0000000000..831501361a --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py @@ -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, + ) diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py new file mode 100644 index 0000000000..aef0805513 --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py @@ -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() diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/utils.py b/tests/v1/e2e/spec_decode/acceptance_rates/utils.py new file mode 100644 index 0000000000..631dc95b81 --- /dev/null +++ b/tests/v1/e2e/spec_decode/acceptance_rates/utils.py @@ -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() diff --git a/tests/v1/e2e/spec_decode/conftest.py b/tests/v1/e2e/spec_decode/conftest.py new file mode 100644 index 0000000000..969450ad96 --- /dev/null +++ b/tests/v1/e2e/spec_decode/conftest.py @@ -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() diff --git a/tests/v1/e2e/spec_decode/draft_model/__init__.py b/tests/v1/e2e/spec_decode/draft_model/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/v1/e2e/spec_decode/test_async_spec_decode.py b/tests/v1/e2e/spec_decode/draft_model/test_async.py similarity index 100% rename from tests/v1/e2e/spec_decode/test_async_spec_decode.py rename to tests/v1/e2e/spec_decode/draft_model/test_async.py diff --git a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py new file mode 100644 index 0000000000..94e32b2375 --- /dev/null +++ b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py @@ -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" + ) diff --git a/tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py b/tests/v1/e2e/spec_decode/draft_model/test_lora.py similarity index 100% rename from tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py rename to tests/v1/e2e/spec_decode/draft_model/test_lora.py diff --git a/tests/v1/e2e/spec_decode/eagle/__init__.py b/tests/v1/e2e/spec_decode/eagle/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py b/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py new file mode 100644 index 0000000000..c5c32d80e5 --- /dev/null +++ b/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py @@ -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, + ) diff --git a/tests/v1/e2e/spec_decode/eagle/utils.py b/tests/v1/e2e/spec_decode/eagle/utils.py new file mode 100644 index 0000000000..998be71af6 --- /dev/null +++ b/tests/v1/e2e/spec_decode/eagle/utils.py @@ -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() diff --git a/tests/v1/e2e/spec_decode/mtp/__init__.py b/tests/v1/e2e/spec_decode/mtp/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/v1/e2e/spec_decode/mtp/test_mtp.py b/tests/v1/e2e/spec_decode/mtp/test_mtp.py new file mode 100644 index 0000000000..45a8236794 --- /dev/null +++ b/tests/v1/e2e/spec_decode/mtp/test_mtp.py @@ -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() diff --git a/tests/v1/e2e/spec_decode/ngram_suffix/__init__.py b/tests/v1/e2e/spec_decode/ngram_suffix/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py b/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py new file mode 100644 index 0000000000..b6dc946a98 --- /dev/null +++ b/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py @@ -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() diff --git a/tests/v1/e2e/spec_decode/speculators/__init__.py b/tests/v1/e2e/spec_decode/speculators/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/tests/v1/e2e/spec_decode/speculators/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/e2e/spec_decode/speculators/test_speculators.py b/tests/v1/e2e/spec_decode/speculators/test_speculators.py new file mode 100644 index 0000000000..27fa21fd76 --- /dev/null +++ b/tests/v1/e2e/spec_decode/speculators/test_speculators.py @@ -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 ` 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." + ) diff --git a/tests/v1/e2e/spec_decode/test_laguna_dflash.py b/tests/v1/e2e/spec_decode/test_laguna_dflash.py deleted file mode 100644 index 1ba9b9749f..0000000000 --- a/tests/v1/e2e/spec_decode/test_laguna_dflash.py +++ /dev/null @@ -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() diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py deleted file mode 100644 index 511f022e58..0000000000 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ /dev/null @@ -1,1679 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os -import random -from collections.abc import Iterable -from dataclasses import dataclass -from typing import Any - -import pytest -import torch -from tqdm import tqdm - -from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline -from tests.utils import ( - get_attn_backend_list_based_on_platform, - large_gpu_mark, - multi_gpu_marks, - multi_gpu_only, - single_gpu_only, - wait_for_rocm_memory_to_settle, -) -from vllm import LLM, SamplingParams -from vllm.assets.base import VLLM_S3_BUCKET_URL -from vllm.assets.image import VLM_IMAGES_DIR -from vllm.benchmarks.datasets import InstructCoderDataset -from vllm.config import CompilationConfig, VllmConfig, replace -from vllm.distributed import cleanup_dist_env_and_memory -from vllm.engine.arg_utils import EngineArgs -from vllm.platforms import current_platform -from vllm.v1.metrics.reader import Metric - -MTP_SIMILARITY_RATE = 0.8 - - -class AsyncSchedulingNotEnabledError(AssertionError): - """Raised when async_scheduling is expected to be True for draft_model - spec decode but is False. Tracked in: - https://github.com/vllm-project/vllm/issues/38929 - """ - - -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]: - 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] - - -@pytest.fixture -def sampling_config(): - return greedy_sampling() - - -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) - - -@pytest.fixture -def model_name(): - return "meta-llama/Llama-3.1-8B-Instruct" - - -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}" - ) - - -@pytest.fixture(autouse=True) -def reset_torch_dynamo(): - """Reset torch dynamo cache before each test""" - yield - # Cleanup after test - torch._dynamo.reset() - - -@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 -@large_gpu_mark(min_gb=20) -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 -@large_gpu_mark(min_gb=20) -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 -@large_gpu_mark(min_gb=20) -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() - - -@pytest.mark.slow_test -@large_gpu_mark(min_gb=80) -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). - - Uses google/gemma-4-12B-it as target with the dspark_gemma4_12b_block7 draft - model. Exercises the full Gemma4 DSpark path (draft build over the reused - base classes DFlashQwen3Model / Qwen3DSparkForCausalLM / Gemma4MTP* / - DSparkMarkovHead, the fused context-KV precompute, the Markov head, and - rejection sampling), so it fails if any base class drifts out of sync. - - 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="google/gemma-4-12B-it", - 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=8192, - max_num_seqs=32, - gpu_memory_utilization=0.8, - 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.588 * 0.9 - assert acceptance_len >= 5.116 * 0.9 - assert gsm8k_accuracy >= 0.937 * 0.9 - finally: - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - -@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 -@large_gpu_mark(min_gb=24) -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 ` 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." - ) - - -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() - - -@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 -@large_gpu_mark(min_gb=24) -@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, - ) - - -@pytest.mark.parametrize( - [ - "model_setup", - "mm_enabled", - "enable_chunked_prefill", - "model_impl", - "expected_accuracy_threshold", - ], - [ - pytest.param( - ( - "eagle", - "meta-llama/Llama-3.1-8B-Instruct", - "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", - 1, - ), - False, - True, - "auto", - 0.7, - marks=large_gpu_mark(min_gb=40), - id="llama3_eagle", - ), - pytest.param( - ( - "eagle", - "meta-llama/Llama-4-Scout-17B-16E-Instruct", - "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", - 4, - ), - False, - False, - "auto", - 0.8, - marks=[*multi_gpu_marks(num_gpus=4), large_gpu_mark(min_gb=40)], - id="llama4_eagle", - ), - pytest.param( - ( - "eagle", - "meta-llama/Llama-4-Scout-17B-16E-Instruct", - "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", - 4, - ), - True, - True, - "auto", - 0.8, - marks=[*multi_gpu_marks(num_gpus=4), large_gpu_mark(min_gb=80)], - id="llama4_eagle_mm", - ), - ], -) -@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) -def test_eagle_correctness_heavy( - 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, - ) - - -@large_gpu_mark(min_gb=24) -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}" - ) - - -@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 -@large_gpu_mark(min_gb=20) -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(MTP_SIMILARITY_RATE * len(ref_outputs)) - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - -@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 - - -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" - ) - - -def get_messages(dataset: str, n: int) -> list[Messages]: - if dataset == "test_prompts": - return get_test_prompts(mm_enabled=False, num_prompts=n) - elif dataset == "likaixin/InstructCoder": - return get_instruct_coder_messages(n=n) - else: - 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 - } - - -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) - - -# Datasets in the format used in DFlash validations -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 - - -@pytest.mark.parametrize( - ["spec_config", "expected_acceptance_lengths", "chat_template_kwargs"], - [ - pytest.param( - dict( - model="Qwen/Qwen3-8B", - trust_remote_code=True, - speculative_config={ - "method": "dflash", - "model": "z-lab/Qwen3-8B-DFlash-b16", - "num_speculative_tokens": 16, - "max_model_len": 32768, - }, - max_model_len=32768, - max_num_seqs=128, - gpu_memory_utilization=0.85, - enforce_eager=False, - disable_log_stats=False, - ), - # All scores from Table 1 in https://arxiv.org/pdf/2602.06036 - { - "mt-bench": 4.24, - "humaneval": 6.50, - # runs with a subset of prompts so extra wide tol here - "gsm8k": 6.54 * 0.975, - }, - {"enable_thinking": False}, - id="dflash", - ), - pytest.param( - dict( - 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, - # Skip multimodal profiling; this is a text-only eval. - limit_mm_per_prompt={"image": 0, "audio": 0}, - disable_log_stats=False, - ), - { - "mt-bench": 2.28, - "humaneval": 2.68, - # runs with a subset of prompts so extra wide tol here - "gsm8k": 2.67 * 0.975, - }, - {}, - id="gemma4", - ), - ], -) -@pytest.mark.parametrize("use_mrv2", [False, True]) -def test_acceptance_rates( - 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() - - -@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, - ) - - -@single_gpu_only -@large_gpu_mark(min_gb=24) -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() - - -@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() - - -@pytest.mark.parametrize("use_mrv2", [False, True]) -def test_dflash_correctness( - monkeypatch: pytest.MonkeyPatch, use_mrv2: bool, dflash_config -): - """ - E2E test for DFlash (block diffusion) speculative decoding. - Ensures output correctness on GSM8k, with cudagraphs and batching on. - """ - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0") - - spec_llm = LLM(**dflash_config) - - # Evaluate GSM8k accuracy (Qwen3-8B ref: ~87-92% on GSM8k) - evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8) - - current_metrics = spec_llm.get_metrics() - acceptance_len = compute_acceptance_len(current_metrics) - - # AR is thoroughly validated in test_dflash_acceptance_rates, in a manner consistent - # with the DFlash paper. However, that test measures AL per-request and thus runs - # with a batch size of 1. To ensure that AL does not collapse with large batch sizes - # we enforce a baseline on the AL over the full lm-eval-style GSM8k test. - expected_len = 3.5 # Measured is 3.9 to 4.0 - print(f"DFlash GSM8k correctness test got AL {acceptance_len}") - assert acceptance_len >= expected_len, ( - "DFlash correctness check failed with" - f" {acceptance_len=}, expected at least {expected_len}" - ) - - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() diff --git a/tests/v1/e2e/spec_decode/utils.py b/tests/v1/e2e/spec_decode/utils.py new file mode 100644 index 0000000000..ccd5232eaa --- /dev/null +++ b/tests/v1/e2e/spec_decode/utils.py @@ -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) diff --git a/vllm/model_executor/models/gemma4_dspark.py b/vllm/model_executor/models/gemma4_dspark.py index 792f108842..a77448da50 100644 --- a/vllm/model_executor/models/gemma4_dspark.py +++ b/vllm/model_executor/models/gemma4_dspark.py @@ -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) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 9ea2ff0f73..13c7644315 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -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(