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