iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Merge branch 'main' into patches

Browse Source
This commit is contained in:
MOVZX
2026-10-17 08:06:56 +00:00
parent be95edb7f6 08a65e287e
commit e0fe5df731
222 changed files with 14823 additions and 3606 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.agents/skills/model-microbenchmark/SKILL.md
+3
View File
@@ -31,6 +31,9 @@ Requests: `[N*]Q[@K][c][:D]`, comma-separated, per DP rank.
| `2048@10240` | prefill after an 8192-token prefix-cache hit |
| `2048@10240c` | chunked prefill: the prefix came from earlier chunks |
Verify cases require `--speculative-config` with `num_speculative_tokens`
at least as large as the case's draft count; insufficient capacity is rejected.
Hit vs chunk only differs on models with SWA bounded replay (DSV4.1): a hit
also recomputes the last window (128 tokens) of the prefix.
.buildkite/test-amd.yaml
+212 -300
View File
@@ -399,27 +399,6 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode/draft_model/
- label: ":amd: (MI250) Spec Decode Speculators + MTP" # TBD
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
dind: false
agent_pool: mi250_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/model_loader/
- vllm/v1/sample/
- vllm/model_executor/layers/
- vllm/transformers_utils/configs/speculators/
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
- vllm/v1/attention/backends/
commands:
- pytest -v -s v1/e2e/spec_decode/speculators/
- pytest -v -s v1/e2e/spec_decode/mtp/
#------------------------------------------------------------- mi250 · misc ------------------------------------------------------------#
- label: ":amd: (MI300) Python-only Installation" # TBD
@@ -1094,25 +1073,6 @@ steps:
#----------------------------------------------------------- mi300 · evals -----------------------------------------------------------#
- label: ":amd: (MI300) LM Eval Small Models" # TBD
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
optional: true
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
- label: ":amd: (MI300) DeepSeek V2-Lite Prefetch Offload Accuracy" # TBD
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1153,23 +1113,6 @@ steps:
commands:
- pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt
- label: ":amd: (MI300) LM Eval KV-Offload" # TBD
timeout_in_minutes: 55
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
- vllm/v1/kv_offload/
- vllm/v1/simple_kv_offload/
- tests/evals/gsm8k/test_gsm8k_offloading.py
- vllm/platforms/rocm.py
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it"
- label: ":amd: (MI300) LM Eval KV-Offload Medium" # TBD
timeout_in_minutes: 70
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1188,24 +1131,6 @@ steps:
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite"
- label: ":amd: (MI300) LM Eval KV-Offload Large" # TBD
timeout_in_minutes: 65
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/distributed/kv_transfer/kv_connector/v1/offloading/
- vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
- vllm/v1/kv_offload/
- vllm/v1/simple_kv_offload/
- tests/evals/gsm8k/test_gsm8k_offloading.py
- vllm/platforms/rocm.py
commands:
- VLLM_ROCM_USE_AITER=1 pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash"
- label: ":amd: (MI300) LM Eval Small Models Harness" # TBD
timeout_in_minutes: 35
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1383,52 +1308,6 @@ steps:
commands:
- VLLM_ENGINE_READY_TIMEOUT_S=1800 bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040
- label: ":amd: (MI300) LM Eval Large Models" # TBD
timeout_in_minutes: 65
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_8
optional: true
num_gpus: 8
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/quantization/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/model_executor/layers/layernorm.py
- csrc/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- tests/evals/
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export PYTORCH_ROCM_ARCH=gfx942
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
- label: ":amd: (MI300) LM Eval Large Models ROCm Harness" # TBD
timeout_in_minutes: 40
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_8
optional: true
num_gpus: 8
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
source_file_dependencies:
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/quantization/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/model_executor/layers/layernorm.py
- csrc/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8 --timeout=5400
- label: ":amd: (MI300) LM Eval Large Models FP8" # TBD
timeout_in_minutes: 40
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1939,20 +1818,6 @@ steps:
commands:
- pytest -v -s models/language/pooling_mteb_test
- label: ":amd: (MI300) Language Models (Extended Pooling) Shard %N"
timeout_in_minutes: 95
parallelism: 4
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/language/pooling
commands:
- pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
- label: ":amd: (MI300) Language Models (Extended Generation)" # TBD
timeout_in_minutes: 105
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1970,21 +1835,6 @@ steps:
#---------------------------------------------------- mi300 · models / multimodal ----------------------------------------------------#
- label: ":amd: (MI300) Multimodal Models (Extended Generation 1)" # TBD
timeout_in_minutes: 100
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/multimodal/generation
- tests/models/multimodal/test_mapping.py
commands:
- pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
- pytest -v -s models/multimodal/test_mapping.py
- label: ":amd: (MI300) Multimodal Models (Extended Generation 2) Shard %N"
timeout_in_minutes: 100
parallelism: 4
@@ -2013,19 +1863,6 @@ steps:
commands:
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
- label: ":amd: (MI300) Multimodal Models (PPL)"
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/multimodal/generation_ppl_test/
commands:
- pytest -v -s models/multimodal/generation_ppl_test/
- label: ":amd: (MI300) Multimodal Models (Standard) 3: llava + qwen2_vl" # TBD
timeout_in_minutes: 55
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2057,19 +1894,6 @@ steps:
commands:
- pytest -v -s models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
- label: ":amd: (MI300) Multimodal Models (Extended Pooling)" # TBD
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/multimodal/pooling
commands:
- pytest -v -s models/multimodal/pooling -m 'not core_model'
#----------------------------------------------------- mi300 · models / quantized -----------------------------------------------------#
- label: ":amd: (MI300) Quantized Models" # TBD
@@ -2533,27 +2357,6 @@ steps:
- VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN
- VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py
- label: ":amd: (MI300) Spec Decode Speculators + MTP" # TBD
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/model_loader/
- vllm/v1/sample/
- vllm/model_executor/layers/
- vllm/transformers_utils/configs/speculators/
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
- vllm/v1/attention/backends/
commands:
- pytest -v -s v1/e2e/spec_decode/speculators/
- pytest -v -s v1/e2e/spec_decode/mtp/
- label: ":amd: (MI300) Spec Decode DeepSeek MTP Parallel Load"
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2581,60 +2384,6 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
- label: ":amd: (MI300) Spec Decode AL DFlash Nightly"
timeout_in_minutes: 120
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/models/qwen3_dflash.py
- vllm/model_executor/models/qwen3_dflash2.py
- vllm/model_executor/models/qwen3_5.py
- vllm/model_executor/models/laguna_dflash.py
- tests/v1/e2e/spec_decode/
- vllm/v1/attention/
- vllm/v1/sample/
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/
- vllm/model_executor/models/qwen2.py
- vllm/model_executor/models/qwen3.py
- vllm/model_executor/models/laguna.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2"
- label: ":amd: (MI300) Spec Decode AL DSpark Nightly"
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
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/
- vllm/v1/attention/
- vllm/v1/sample/
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/
- vllm/model_executor/models/qwen2.py
- vllm/model_executor/models/qwen3.py
- vllm/model_executor/models/qwen3_dflash.py
- vllm/model_executor/models/gemma4.py
- vllm/model_executor/models/gemma4_mm.py
- vllm/model_executor/models/gemma4_unified.py
- vllm/model_executor/models/gemma4_mtp.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/
- label: ":amd: (MI300) V1 Spec Decode" # TBD
timeout_in_minutes: 50
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2719,39 +2468,6 @@ steps:
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
- pytest -v -s v1/test_tensor_ipc_queue.py
- label: ":amd: (MI300) Speculators Correctness Nightly"
timeout_in_minutes: 40
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/config/speculative.py
- vllm/engine/arg_utils.py
- vllm/transformers_utils/config.py
- vllm/transformers_utils/configs/speculators/
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/v1/worker/gpu_model_runner.py
- vllm/v1/sample/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/
- vllm/model_executor/models/llama_eagle3.py
- vllm/model_executor/models/qwen3.py
- vllm/model_executor/models/qwen3_dflash.py
- vllm/model_executor/models/registry.py
- vllm/_aiter_ops.py
- tests/evals/gsm8k/
- tests/v1/spec_decode/test_speculators_correctness.py
- vllm/platforms/rocm.py
- vllm/v1/worker/gpu/spec_decode/dflash/
commands:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
- label: ":amd: (MI300) Extract Hidden States Integration" # TBD
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -3106,22 +2822,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: ":amd: (MI300) V1 E2E Hybrid Chunked Prefill" # TBD
timeout_in_minutes: 35
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_4
num_gpus: 4
optional: true
source_file_dependencies:
- vllm/v1/attention/backends/utils.py
- vllm/v1/worker/gpu_model_runner.py
- tests/v1/e2e/test_hybrid_chunked_prefill.py
- vllm/v1/core/
working_dir: /vllm-workspace/tests
commands:
- pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py
#########################################################################################################################################
# #
# MI355 (gfx950) tests #
@@ -4604,6 +4304,82 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi355-nemotron-3-super.txt
- label: ":amd: (MI355) LM Eval Small Models" # TBD
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_dpx
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
optional: true
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
- label: ":amd: (MI355 DPX) LM Eval KV-Offload" # TBD
timeout_in_minutes: 55
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_dpx
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
- vllm/v1/kv_offload/
- vllm/v1/simple_kv_offload/
- tests/evals/gsm8k/test_gsm8k_offloading.py
- vllm/platforms/rocm.py
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it"
- label: ":amd: (MI355) LM Eval KV-Offload Large" # TBD
timeout_in_minutes: 65
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_4
num_gpus: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
- vllm/v1/kv_offload/
- vllm/v1/simple_kv_offload/
- tests/evals/gsm8k/test_gsm8k_offloading.py
- vllm/platforms/rocm.py
commands:
- VLLM_ROCM_USE_AITER=1 pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash"
- label: ":amd: (MI355) LM Eval Large Models ROCm Harness" # TBD
timeout_in_minutes: 40
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_8
optional: true
num_gpus: 8
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
source_file_dependencies:
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/quantization/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/model_executor/layers/layernorm.py
- csrc/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8 --timeout=5400
#--------------------------------------------------------- mi355 · examples ----------------------------------------------------------#
- label: ":amd: (MI355 DPX) Examples" # TBD
@@ -5421,6 +5197,19 @@ steps:
- pytest models/multimodal/generation/test_memory_leak.py -m core_model
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model
- label: ":amd: (MI355 DPX) Multimodal Models (PPL)"
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_dpx
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/models/multimodal/generation_ppl_test/
commands:
- pytest -v -s models/multimodal/generation_ppl_test/
#----------------------------------------------------- mi355 · models / quantized -----------------------------------------------------#
- label: ":amd: (MI355 DPX) Quantized Models" # TBD
@@ -5565,6 +5354,27 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode/draft_model/
- label: ":amd: (MI355 DPX) Spec Decode Speculators + MTP" # TBD
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_dpx
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/model_loader/
- vllm/v1/sample/
- vllm/model_executor/layers/
- vllm/transformers_utils/configs/speculators/
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
- vllm/v1/attention/backends/
commands:
- pytest -v -s v1/e2e/spec_decode/speculators/
- pytest -v -s v1/e2e/spec_decode/mtp/
- label: ":amd: (MI355 DPX) Spec Decode Eagle 1: DeepSeek + Qwen"
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -6041,6 +5851,108 @@ steps:
- if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]"; fi
- if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi
- label: ":amd: (MI355 DPX) V1 E2E Hybrid Chunked Prefill" # TBD
timeout_in_minutes: 35
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_dpx
optional: true
source_file_dependencies:
- vllm/v1/attention/backends/utils.py
- vllm/v1/worker/gpu_model_runner.py
- tests/v1/e2e/test_hybrid_chunked_prefill.py
- vllm/v1/core/
working_dir: /vllm-workspace/tests
commands:
- pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py
- label: ":amd: (MI355 DPX) Speculators Correctness Nightly"
timeout_in_minutes: 40
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_dpx
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/config/speculative.py
- vllm/engine/arg_utils.py
- vllm/transformers_utils/config.py
- vllm/transformers_utils/configs/speculators/
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/v1/worker/gpu_model_runner.py
- vllm/v1/sample/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/
- vllm/model_executor/models/llama_eagle3.py
- vllm/model_executor/models/qwen3.py
- vllm/model_executor/models/qwen3_dflash.py
- vllm/model_executor/models/registry.py
- vllm/_aiter_ops.py
- tests/evals/gsm8k/
- tests/v1/spec_decode/test_speculators_correctness.py
- vllm/platforms/rocm.py
- vllm/v1/worker/gpu/spec_decode/dflash/
commands:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
- label: ":amd: (MI355 DPX) Spec Decode AL DFlash Nightly"
timeout_in_minutes: 120
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_dpx
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/models/qwen3_dflash.py
- vllm/model_executor/models/qwen3_dflash2.py
- vllm/model_executor/models/qwen3_5.py
- vllm/model_executor/models/laguna_dflash.py
- tests/v1/e2e/spec_decode/
- vllm/v1/attention/
- vllm/v1/sample/
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/
- vllm/model_executor/models/qwen2.py
- vllm/model_executor/models/qwen3.py
- vllm/model_executor/models/laguna.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2"
- label: ":amd: (MI355 DPX) Spec Decode AL DSpark Nightly"
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_dpx
optional: true
working_dir: "/vllm-workspace/tests"
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/
- vllm/v1/attention/
- vllm/v1/sample/
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/
- vllm/model_executor/models/qwen2.py
- vllm/model_executor/models/qwen3.py
- vllm/model_executor/models/qwen3_dflash.py
- vllm/model_executor/models/gemma4.py
- vllm/model_executor/models/gemma4_mm.py
- vllm/model_executor/models/gemma4_unified.py
- vllm/model_executor/models/gemma4_mtp.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/
#----------------------------------------------------------- mi355 · misc ------------------------------------------------------------#
- label: ":amd: (MI355 DPX) Regression" # TBD
.buildkite/test_areas/engine.yaml
+2 -2
View File
@@ -238,9 +238,9 @@ steps:
- pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py
mirror:
amd:
label: ":amd: (MI300) V1 E2E Hybrid Chunked Prefill"
label: ":amd: (MI355 DPX) V1 E2E Hybrid Chunked Prefill"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 35
depends_on:
- image-build-amd
.buildkite/test_areas/kernels.yaml
+1 -2
View File
@@ -545,7 +545,6 @@ steps:
num_devices: 1
source_file_dependencies:
- cmake/external_projects/deepgemm.cmake
- tools/install_deepgemm.sh
- tools/build_deepgemm_C.py
- tools/setup_deepgemm_pythons.sh
- tools/check_wheel_deepgemm.py
@@ -939,7 +938,7 @@ steps:
- label: ":nvidia: (B200) FusedMoE Layer Kernels"
key: kernels-fusedmoe-layer-test-2-b200s
timeout_in_minutes: 90
timeout_in_minutes: 45
device: b200-k8s
num_devices: 2
source_file_dependencies:
.buildkite/test_areas/lm_eval.yaml
+11 -12
View File
@@ -15,9 +15,9 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
mirror:
amd:
label: ":amd: (MI300) LM Eval Small Models"
label: ":amd: (MI355) LM Eval Small Models"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 45
depends_on:
- image-build-amd
@@ -707,18 +707,17 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi355-nemotron-3-super.txt
- label: ":amd: (MI300) LM Eval Large Models"
- label: ":amd: (MI355) LM Eval Large Models"
key: amd-lm-eval-large-models
dind: false
device: mi300_8
device: mi355_8
optional: true
timeout_in_minutes: 65
depends_on:
- image-build-amd
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-gfx950-large.txt
source_file_dependencies:
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
@@ -1298,9 +1297,9 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it"
mirror:
amd:
label: ":amd: (MI300) LM Eval KV-Offload"
label: ":amd: (MI355 DPX) LM Eval KV-Offload"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 55
depends_on:
- image-build-amd
@@ -1360,9 +1359,9 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash"
mirror:
amd:
label: ":amd: (MI300) LM Eval KV-Offload Large"
label: ":amd: (MI355) LM Eval KV-Offload Large"
dind: false
device: mi300_4
device: mi355_4
timeout_in_minutes: 65
depends_on:
- image-build-amd
@@ -1427,10 +1426,10 @@ steps:
# Qwen2-57B has 28 attention heads and cannot run at TP8.
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
- label: ":amd: (MI300) LM Eval Large Models ROCm Harness"
- label: ":amd: (MI355) LM Eval Large Models ROCm Harness"
key: amd-lm-eval-large-models-rocm-harness
dind: false
device: mi300_8
device: mi355_8
optional: true
timeout_in_minutes: 40
depends_on:
.buildkite/test_areas/misc.yaml
+54
View File
@@ -164,6 +164,59 @@ steps:
depends_on:
- image-build-amd
# The shared KV Connectors steps above exclude vllm/distributed/kv_transfer/ from
# their triggers, and CUDA skips MoRIIO's mori/aiter-gated tests. Run the whole
# MoRIIO suite on ROCm whenever the connector or its tests change.
- label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests"
key: amd-moriio-connector-unit-mi355
dind: false
device: mi355_dpx
num_devices: 1
timeout_in_minutes: 30
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/moriio/
- vllm/distributed/kv_transfer/kv_connector/v1/base.py
- vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
- examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
- tests/v1/kv_connector/unit/test_moriio_
- tests/v1/kv_connector/unit/utils.py
- requirements/kv_connectors_rocm.txt
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s -rs v1/kv_connector/unit/test_moriio_*.py
# Single-node MoRIIO WRITE P/D on dummy weights: prefill and decode on separate
# GPUs over xGMI, compared exactly against a direct serve.
- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N"
key: amd-moriio-write-dummy-e2e-mi355
dind: false
device: mi355_2
num_devices: 2
# ~3.5-5 min per case; pytest-shard hashes test ids, so with 20 cases the
# largest of 6 shards holds 5 (~22 min on MI355). Every case picks free ports
# for its servers and proxy and its own compile-cache dir.
parallelism: 6
timeout_in_minutes: 45
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/moriio/
- vllm/distributed/kv_transfer/kv_connector/v1/base.py
- vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
- vllm/model_executor/model_loader/dummy_loader.py
- examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
- tests/v1/kv_connector/moriio_integration/
- requirements/kv_connectors_rocm.txt
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- label: ":nvidia: (H200 MIG 35GB) V1 Metrics + LM Eval"
device: h200_35gb
key: v1-metrics-lmeval
@@ -214,6 +267,7 @@ steps:
- vllm/v1/
- tests/v1
- tests/watermarking
- examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py
device: cpu-small
commands:
# split the test to avoid interference
.buildkite/test_areas/models_language.yaml
+2 -2
View File
@@ -216,9 +216,9 @@ steps:
parallelism: 4
commands:
- "pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
label: ":amd: (MI300) Language Models (Extended Pooling) Shard %N"
label: ":amd: (MI355 DPX) Language Models (Extended Pooling) Shard %N"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 95
depends_on:
- image-build-amd
.buildkite/test_areas/models_multimodal.yaml
+6 -6
View File
@@ -163,9 +163,9 @@ steps:
- pytest -v -s models/multimodal/test_mapping.py
mirror:
amd:
label: ":amd: (MI300) Multimodal Models (Extended Generation 1)"
label: ":amd: (MI355 DPX) Multimodal Models (Extended Generation 1)"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 100
depends_on:
- image-build-amd
@@ -182,9 +182,9 @@ steps:
- pytest -v -s models/multimodal/generation_ppl_test/
mirror:
amd:
label: ":amd: (MI300) Multimodal Models (PPL)"
label: ":amd: (MI355 DPX) Multimodal Models (PPL)"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 90
depends_on:
- image-build-amd
@@ -240,9 +240,9 @@ steps:
- pytest -v -s models/multimodal/pooling -m 'not core_model'
mirror:
amd:
label: ":amd: (MI300) Multimodal Models (Extended Pooling)"
label: ":amd: (MI355 DPX) Multimodal Models (Extended Pooling)"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 75
depends_on:
- image-build-amd
.buildkite/test_areas/spec_decode.yaml
+14 -14
View File
@@ -137,9 +137,9 @@ steps:
- pytest -v -s v1/e2e/spec_decode/speculators/
mirror:
amd:
label: ":amd: (MI300) Spec Decode Speculators"
label: ":amd: (MI355 DPX) Spec Decode Speculators"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 40
depends_on:
- image-build-amd
@@ -167,9 +167,9 @@ steps:
- pytest -v -s v1/e2e/spec_decode/mtp/gemma4/
mirror:
amd:
label: ":amd: (MI300) Spec Decode MTP Gemma4"
label: ":amd: (MI355 DPX) Spec Decode MTP Gemma4"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 40
depends_on:
- image-build-amd
@@ -198,9 +198,9 @@ steps:
- pytest -v -s v1/e2e/spec_decode/mtp/mimo/
mirror:
amd:
label: ":amd: (MI300) Spec Decode MTP DeepSeek + MiMo"
label: ":amd: (MI355 DPX) Spec Decode MTP DeepSeek + MiMo"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 40
depends_on:
- image-build-amd
@@ -228,9 +228,9 @@ steps:
- pytest -v -s v1/e2e/spec_decode/mtp/qwen3_5/
mirror:
amd:
label: ":amd: (MI300) Spec Decode MTP Qwen3.5"
label: ":amd: (MI355 DPX) Spec Decode MTP Qwen3.5"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 40
depends_on:
- image-build-amd
@@ -387,9 +387,9 @@ steps:
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
mirror:
amd:
label: ":amd: (MI300) Speculators Correctness Nightly"
label: ":amd: (MI355 DPX) Speculators Correctness Nightly"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 40
depends_on:
- image-build-amd
@@ -452,9 +452,9 @@ steps:
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2"
mirror:
amd:
label: ":amd: (MI300) Spec Decode AL DFlash Nightly"
label: ":amd: (MI355 DPX) Spec Decode AL DFlash Nightly"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 120
commands:
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2"
@@ -531,9 +531,9 @@ steps:
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/
mirror:
amd:
label: ":amd: (MI300) Spec Decode AL DSpark Nightly"
label: ":amd: (MI355 DPX) Spec Decode AL DSpark Nightly"
dind: false
device: mi300_1
device: mi355_dpx
timeout_in_minutes: 90
depends_on:
- image-build-amd
.github/CODEOWNERS
+2 -1
View File
@@ -8,7 +8,7 @@
/vllm/distributed/kv_transfer/kv_connector/v1/hisparse @MatthewBonanni @LucasWilkinson
/vllm/lora @jeejeelee
/vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye @wzhao18
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye @wzhao18 @bnellnm
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
/vllm/model_executor/layers/mamba @tdoublep @tomeras91 @ZJY0516
/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
@@ -138,6 +138,7 @@
/tests/evals @mgoin @vadiklyutiy @AndreasKaratzas @ZJY0516
/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas
/tests/kernels/ir @ProExpertProg @tjtanaa
/tests/kernels/moe @bnellnm
/tests/models @DarkLight1337 @ywang96 @AndreasKaratzas
/tests/multimodal @DarkLight1337 @ywang96 @NickLucche @shen-shanshan
/tests/multimodal/media @Isotr0py
cmake/external_projects/deepgemm.cmake
-1
View File
@@ -27,7 +27,6 @@ if(DEEPGEMM_SRC_DIR)
set(deepgemm_SOURCE_DIR "${DEEPGEMM_SRC_DIR}")
message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
else()
# Keep in sync with tools/install_deepgemm.sh
set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git")
# Pinned to the fork's dev branch: upstream's 26/09/30 release (2.8.1)
# plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port,
csrc/activation_kernels.cu
-27
View File
@@ -1059,33 +1059,6 @@ void situ_and_mul(torch::stable::Tensor& out, // [..., d]
});
}
void masked_situ_and_mul(torch::stable::Tensor& out, // [E, T, d]
torch::stable::Tensor& input, // [E, T, 2 * d]
const torch::stable::Tensor& expert_num_tokens,
double beta, double linear_beta) {
int num_experts = input.size(0);
int max_num_tokens = input.size(1);
int d = input.size(2) / 2;
if (num_experts == 0 || max_num_tokens == 0 || d == 0) {
return;
}
constexpr int block_size = 256;
dim3 grid((d + block_size - 1) / block_size, num_experts);
dim3 block(block_size);
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
const cudaStream_t stream = get_current_cuda_stream();
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
input.scalar_type(), "masked_situ_and_mul_kernel", [&] {
vllm::masked_situ_and_mul_kernel<scalar_t, true>
<<<grid, block, 0, stream>>>(
out.mutable_data_ptr<scalar_t>(),
input.const_data_ptr<scalar_t>(),
expert_num_tokens.const_data_ptr<int>(), max_num_tokens, d,
(float)beta, (float)linear_beta);
});
}
#define LAUNCH_MASKED_ACT_AND_MUL(KERNEL, ACT_FIRST, HAS_CLAMP, LIMIT, ALPHA, \
BETA) \
vllm::masked_act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST, \
csrc/fp32_router_gemm.cu
+76 -5
View File
@@ -1,7 +1,7 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=32, for the
// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=48, for the
// supported (E, H) pairs listed at the bottom of this file.
// Supports bf16 or fp32 activation; weight is always fp32.
@@ -214,8 +214,26 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
// production-fidelity harness (CUDA-graph replay, per-layer cold weights).
// Only enabled on the Blackwell family where it was validated; Hopper and
// other shapes / fp32 activation keep the legacy geometry.
if constexpr (std::is_same_v<InputT, __nv_bfloat16> && kNumExperts == 128 &&
kHiddenDim == 6144) {
if constexpr (std::is_same_v<InputT, __nv_bfloat16> && kNumExperts == 256 &&
kHiddenDim == 5120) {
if (!isBlackwellFamily()) {
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
return;
}
if constexpr (kNumTokens > 16 && kNumTokens % 8 == 0) {
constexpr int block_size = kNumTokens == 24 ? 160 : 128;
launchFp32RouterGemm<InputT, block_size, 2, kNumTokens, kNumExperts,
kHiddenDim, 4>(output, mat_a, mat_b, stream);
} else {
constexpr int block_size = kNumTokens == 1 ? 64 : 128;
constexpr int token_groups = kNumTokens % 2 == 0 ? 2 : 1;
launchFp32RouterGemm<InputT, block_size, 1, kNumTokens, kNumExperts,
kHiddenDim, token_groups>(output, mat_a, mat_b,
stream);
}
} else if constexpr (std::is_same_v<InputT, __nv_bfloat16> &&
kNumExperts == 128 && kHiddenDim == 6144) {
// MiniMax-M3. Legacy 128/1 only fills 128 blocks and pays the same
// accumulator register cliffs; B300 sweep:
// even M in [6, 10] : BS=384, EPB=1, 2 token groups (1.26-1.43x)
@@ -261,6 +279,27 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
}
} else if constexpr (std::is_same_v<InputT, __nv_bfloat16> &&
kNumExperts == 256 && kHiddenDim == 8192) {
// MuseSpark1x: the 8MB weight makes this shape FFMA-bound past M~10, so
// the wider block wins at every M and token groups clear the M >= 28
// accumulator cliff. Thresholds are from a B300 sweep. BS=192/384 are not
// instantiable: VPT=8 requires kBlockSize to divide kHiddenDim / 8 = 1024.
if (!isBlackwellFamily()) {
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
return;
}
if constexpr (kNumTokens >= 10 && kNumTokens % 2 == 0) {
launchFp32RouterGemm<InputT, 256, 2, kNumTokens, kNumExperts, kHiddenDim,
2>(output, mat_a, mat_b, stream);
} else if constexpr (kNumTokens >= 13) {
launchFp32RouterGemm<InputT, 256, 2, kNumTokens, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
} else {
launchFp32RouterGemm<InputT, 256, 1, kNumTokens, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
}
} else {
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
@@ -268,8 +307,9 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
}
// ---------------------------------------------------------------------------
// Explicit instantiations: M=1..32, for both input types, for the supported
// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3].
// Explicit instantiations for both input types.
// Supported (E, H) pairs: (256, 3072), (128, 6144), (256, 5120),
// (352, 4096) [Step-5], and (256, 8192) [Muse Spark].
// ---------------------------------------------------------------------------
#define INSTANTIATE(T, M, E, H) \
@@ -315,5 +355,36 @@ INSTANTIATE_ALL(__nv_bfloat16, 256, 3072)
INSTANTIATE_ALL(float, 128, 6144)
INSTANTIATE_ALL(__nv_bfloat16, 128, 6144)
INSTANTIATE_ALL(float, 256, 5120)
INSTANTIATE_ALL(__nv_bfloat16, 256, 5120)
INSTANTIATE_ALL(float, 352, 4096)
INSTANTIATE_ALL(__nv_bfloat16, 352, 4096)
INSTANTIATE_ALL(float, 256, 8192)
INSTANTIATE_ALL(__nv_bfloat16, 256, 8192)
#define INSTANTIATE_LARGE(T, E, H) \
INSTANTIATE(T, 33, E, H) \
INSTANTIATE(T, 34, E, H) \
INSTANTIATE(T, 35, E, H) \
INSTANTIATE(T, 36, E, H) \
INSTANTIATE(T, 37, E, H) \
INSTANTIATE(T, 38, E, H) \
INSTANTIATE(T, 39, E, H) \
INSTANTIATE(T, 40, E, H) \
INSTANTIATE(T, 41, E, H) \
INSTANTIATE(T, 42, E, H) \
INSTANTIATE(T, 43, E, H) \
INSTANTIATE(T, 44, E, H) \
INSTANTIATE(T, 45, E, H) \
INSTANTIATE(T, 46, E, H) \
INSTANTIATE(T, 47, E, H) \
INSTANTIATE(T, 48, E, H)
INSTANTIATE_LARGE(float, 256, 5120)
INSTANTIATE_LARGE(__nv_bfloat16, 256, 5120)
#undef INSTANTIATE_LARGE
#undef INSTANTIATE_ALL
#undef INSTANTIATE
csrc/fp32_router_gemm_entry.cu
+20 -7
View File
@@ -25,10 +25,13 @@ inline int getSMVersion() {
static constexpr int FP32_MAX_TOKENS = 32;
// Supported (hidden_dim, num_experts) pairs (must match the instantiations in
// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3.
// fp32_router_gemm.cu).
static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) {
return (hidden_dim == 3072 && num_experts == 256) ||
(hidden_dim == 6144 && num_experts == 128);
(hidden_dim == 6144 && num_experts == 128) ||
(hidden_dim == 5120 && num_experts == 256) ||
(hidden_dim == 4096 && num_experts == 352) ||
(hidden_dim == 8192 && num_experts == 256);
}
// Forward declarations — 4 template params must match fp32_router_gemm.cu
@@ -60,8 +63,7 @@ struct Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kEnd, kEnd> {
invokeFp32RouterGemm<InputT, kEnd, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
} else {
throw std::invalid_argument(
"fp32_router_gemm: num_tokens must be in [1, 32]");
throw std::invalid_argument("fp32_router_gemm: unsupported num_tokens");
}
}
};
@@ -77,6 +79,15 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
} else if (num_experts == 128 && hidden_dim == 6144) {
Fp32LoopUnroller<InputT, 128, 6144, 1, FP32_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, stream);
} else if (num_experts == 256 && hidden_dim == 5120) {
Fp32LoopUnroller<InputT, 256, 5120, 1, 48>::unroll(num_tokens, output,
mat_a, mat_b, stream);
} else if (num_experts == 352 && hidden_dim == 4096) {
Fp32LoopUnroller<InputT, 352, 4096, 1, FP32_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, stream);
} else if (num_experts == 256 && hidden_dim == 8192) {
Fp32LoopUnroller<InputT, 256, 8192, 1, FP32_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, stream);
} else {
throw std::invalid_argument(
"fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
@@ -111,9 +122,11 @@ void fp32_router_gemm(
STD_TORCH_CHECK(
fp32_router_gemm_supported(hidden_dim, num_experts),
"fp32_router_gemm: supported (hidden_dim, num_experts) pairs are "
"(3072, 256) and (6144, 128)");
STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
"fp32_router_gemm: num_tokens must be in [0, 32]");
"(3072, 256), (6144, 128), (5120, 256), (4096, 352), and (8192, 256)");
const int max_tokens =
hidden_dim == 5120 && num_experts == 256 ? 48 : FP32_MAX_TOKENS;
STD_TORCH_CHECK(num_tokens <= max_tokens,
"fp32_router_gemm: num_tokens exceeds supported range");
STD_TORCH_CHECK(
mat_a.scalar_type() == torch::headeronly::ScalarType::Float ||
mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
csrc/ops.h
-4
View File
@@ -579,10 +579,6 @@ void situ_and_mul_quant(
int64_t group_size = 0,
std::optional<torch::stable::Tensor> num_valid_tokens = std::nullopt,
int64_t topk = 1);
void masked_situ_and_mul(torch::stable::Tensor& out,
torch::stable::Tensor& input,
const torch::stable::Tensor& expert_num_tokens,
double beta = 1.0, double linear_beta = -1.0);
void masked_moe_activation(torch::stable::Tensor& out,
torch::stable::Tensor& input,
const torch::stable::Tensor& valid_token_counts,
csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu
+7
View File
@@ -135,6 +135,13 @@ hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) {
#pragma unroll
for (int64_t k = 0; k < num_chunks; k++) {
// The last warp may own fewer chunks than num_chunks: loading the
// rest would read past the end of `a`. The cp.async.wait_group counts
// below already assume only real_num_chunks groups were committed.
if constexpr(enable_mask) {
if (k >= real_num_chunks)
break;
}
size_t shared_ptr = __cvta_generic_to_shared(b_frag_ptr);
#if (__CUDA_ARCH__ >= 900) // SM90
asm volatile(
csrc/torch_bindings.cpp
-4
View File
@@ -640,9 +640,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"situ_and_mul_quant(Tensor! out, Tensor! scale, Tensor input, "
"float beta=1.0, float linear_beta=-1.0, int group_size=0, "
"Tensor? num_valid_tokens=None, int topk=1) -> ()");
ops.def(
"masked_situ_and_mul(Tensor! out, Tensor input, Tensor "
"expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()");
ops.def(
"masked_moe_activation(Tensor! out, Tensor input, Tensor "
"valid_token_counts, str activation, float clamp_limit=0.0, float "
@@ -875,7 +872,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul));
ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul));
ops.impl("situ_and_mul_quant", TORCH_BOX(&situ_and_mul_quant));
ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul));
ops.impl("masked_moe_activation", TORCH_BOX(&masked_moe_activation));
ops.impl("gelu_new", TORCH_BOX(&gelu_new));
ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
docs/features/kv_offloading_usage.md
+29
View File
@@ -411,6 +411,35 @@ Example (OpenAI-compatible completions request):
}
```
## Explicit pool sizing for `SimpleCPUOffloadConnector`
`SimpleCPUOffloadConnector` accepts an optional `num_cpu_blocks` in
`kv_connector_extra_config`. The scheduler and every worker use this exact
number of offload blocks instead of automatic sizing from
`cpu_bytes_to_use_per_rank` or `cpu_bytes_to_use / world_size`.
Use the same positive integer on all ranks, chosen
to fit the smallest per-worker capacity. This supports pipeline stages with
different KV cache sizes without adding a collective during initialization.
```json
{
"kv_connector": "SimpleCPUOffloadConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"cpu_bytes_to_use_per_rank": 1073741824,
"num_cpu_blocks": 64,
"lazy_offload": false
}
}
```
The example count must be adjusted for the model, KV layout, and parallelism.
The configured byte budget remains an upper bound: startup fails if the
requested count exceeds a worker's capacity. For the `disk` backend the count
selects disk slots and is checked against `disk_capacity_bytes`. Omitting
`num_cpu_blocks` preserves automatic per-worker sizing. This option applies
to `SimpleCPUOffloadConnector`, separately from `OffloadingConnector` above.
## Further Reading
- [vLLM blog: KV Offloading Connector](https://vllm.ai/blog/2026-01-08-kv-offloading-connector) — motivation, architecture (DMA-based async transfer), and benchmarks (TTFT and throughput).
docs/features/moriio_connector_usage.md
+8
View File
@@ -136,6 +136,14 @@ WRITE mode is used by default. READ mode can be configured by setting `--kv-tran
!!! note
`notify_port` is used as a *base* port: each (DP rank, TP rank) pair within an instance uses `notify_port + offset` where the offset is based on the rank. Make sure the range starting at `notify_port` is free on the host.
### Hybrid attention and recurrent state
WRITE mode supports hybrid attention plus Mamba/KDA models, including DSpark, with equal producer and consumer TP sizes. Use matching cache layouts on both instances, `--mamba-cache-mode align`, and `read_mode: false` in each instance's `kv_connector_extra_config`. Both peers must support hybrid WRITE. Sliding-window WRITE and hybrid speculative methods other than DSpark remain unsupported.
For DSpark, configure the same draft model on both instances through `--speculative-config`. Target and draft attention groups retain their own block IDs and token spans; only the running recurrent state is transferred, not speculative scratch slots. The producer transfers the state after the first N-1 prompt tokens, and the consumer recomputes the last prompt token. Convolution and SSM transfers must both finish before the layer is complete.
Hybrid READ keeps its existing single-transferable-attention-group requirement. Plain-attention WRITE retains its existing notification format.
### Transport configuration
MoRI has two transport backends: RDMA and xGMI. You can select backend using `--kv-transfer-config.kv_connector_extra_config.backend $BACKEND`, with `$BACKEND` being `rdma` or `xgmi`. RDMA is the default backend and should be used in multi-node deployments.
docs/features/multimodal_inputs.md
+3 -3
View File
@@ -187,7 +187,7 @@ Multi-image input can be extended to perform video captioning. We show this with
#### Custom RGBA Background Color
When loading RGBA images (images with transparency), vLLM converts them to RGB format. By default, transparent pixels are replaced with white background. You can customize this background color using the `rgba_background_color` parameter in `media_io_kwargs`.
When loading images with transparency (including RGBA images and PNGs with a tRNS chunk), vLLM converts them to RGB format. By default, transparent pixels are replaced with white background. You can customize this background color using the `rgba_background_color` parameter in `media_io_kwargs`.
??? code
@@ -212,7 +212,7 @@ When loading RGBA images (images with transparency), vLLM converts them to RGB f
!!! note
- The `rgba_background_color` accepts RGB values as a list `[R, G, B]` or tuple `(R, G, B)` where each value is 0-255
- This setting only affects RGBA images with transparency; RGB images are unchanged
- This setting affects images with transparency, including RGB PNGs with a transparent color; opaque RGB images are unchanged
- If not specified, the default white background `(255, 255, 255)` is used for backward compatibility
#### Moondream3 Prompt Recipes { #moondream3-prompt-recipes }
@@ -1089,7 +1089,7 @@ By passing this metadata, the model can better understand the temporal distribut
#### Custom RGBA Background Color
To use a custom background color for RGBA images, pass the `rgba_background_color` parameter via `--media-io-kwargs`:
To use a custom background color for images with transparency, pass the `rgba_background_color` parameter via `--media-io-kwargs`:
```bash
# Example: Black background for dark theme
docs/features/tool_calling.md
+24
View File
@@ -201,6 +201,30 @@ checkpoint's output format.
!!! note
With `tool_choice="auto"`, structural-tag constraints require both `VLLM_ENFORCE_STRICT_TOOL_CALLING=true` (the default) and at least one tool with `strict: true`, or a server-side floor set via `--tool-strict-level`. When these conditions are met and the selected parser supports structural tags, vLLM constrains the tool-call envelope and pins the argument schema of each tool that sets `strict: true` (or of every tool under `--tool-strict-level parameter`). Otherwise, vLLM extracts tool calls from raw text, so arguments may occasionally be malformed or violate the function's parameter schema.
### Beam Models (`beam`)
Beam checkpoints describe their output format with `response_template` metadata, so the `beam` parser reads it like the `hf` parser does:
```bash
vllm serve <model> \
--enable-auto-tool-choice \
--tool-call-parser beam \
--reasoning-parser beam
```
Unlike `hf`, the `beam` parser also constrains its output. For strict tools, required or named tool choice, and structured outputs, it derives a
structural-tag grammar from the template. The grammar follows the template's tool call format, pins the arguments of strict tools when their
encoding can be derived, allows a single call with `parallel_tool_calls=false`, and frames a structured output as the answer. Raw string arguments
keep the `format` or `pattern` of their schema, as XGrammar keeps them in JSON strings, unless a match could end the value early. With
`tool_choice="none"`, the grammar allows only reasoning and an answer, so a model that still sees the tools in its prompt cannot call them. When no
grammar can be derived, required and named tool choice are rejected, and strict tools under `tool_choice="auto"` are left unconstrained with a
warning.
It also constrains the output of requests without strict tools, forced tool choice, or structured outputs to the template's format: tool calls with
free arguments or an answer, or just the rest of an answer the prompt already opened. This always-on grammar needs the `xgrammar` structured
outputs backend, the default, and is disabled with `VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT=0`. From the message start, every `beam` grammar opens with
reasoning that holds visible text, unless `enable_thinking` is false.
### Hermes Models (`hermes`)
All Nous Research Hermes-series models newer than Hermes 2 Pro should be supported.
docs/models/pooling_models/scoring.md
+2
View File
@@ -455,6 +455,8 @@ Score templates are supported for **cross-encoder** models only. If you are usin
Some scoring models require a specific prompt format to work correctly. You can specify a custom score template using the `--chat-template` parameter (see [Chat Template](../../serving/online_serving/README.md#chat-template)).
For offline scoring, pass `chat_template` to the `LLM` constructor to use it for every `LLM.score` call. A `chat_template` passed to `LLM.score` overrides the constructor's template for that call.
Like chat templates, the score template receives a `messages` list. For scoring, each message has a `role` attribute—either `"query"` or `"document"`. For the usual kind of point-wise cross-encoder, you can expect exactly two messages: one query and one document. To access the query and document content, use Jinja's `selectattr` filter:
- **Query**: `{{ (messages | selectattr("role", "eq", "query") | first).content }}`
docs/serving/online_serving/derenderer.md
+2 -2
View File
@@ -105,7 +105,7 @@ The server keeps no state between calls. Everything the next call needs is in `s
- **Send `prompt_token_ids` on the first chunk, even without a parser.** It seeds detokenization so the first token's leading space matches the coupled endpoint (see [Request format](#request-format)). Later chunks already carry that context in `stream_state`.
- **Don't forward `[DONE]`.** It marks the end of the generate stream and isn't a chunk.
`/inference/v1/generate` returns `GenerateLogProbs` with integer `token_id`s (the generate server has no tokenizer), and derender turns those into `ChatCompletionLogProbs` / `CompletionLogProbs`, filling `token` and `bytes` from the tokenizer with the usual U+FFFD byte-fallback correction. Generate returns every top-k candidate the engine produced; derender keeps the ones the coupled endpoint would, with the same helpers: the first `top_logprobs` on chat (all of them for `top_logprobs=-1`) and the first `logprobs + 1` on completions (all of them for `logprobs=-1`). It reads those values from `chat_request` / `completion_request`, so without the original request, or with one that does not ask for logprobs, it returns every candidate. Requests that use `logprob_token_ids` get no output logprobs from generate (their `sampling_params.logprobs` is unset), so derender has none to return for them either. Streaming derender does the same per chunk on the plain detokenization path (see [Streaming state and logprobs](#streaming-state-and-logprobs)). The parser path doesn't resolve them yet and drops them.
`/inference/v1/generate` returns `GenerateLogProbs` with integer `token_id`s (the generate server has no tokenizer), and derender turns those into `ChatCompletionLogProbs` / `CompletionLogProbs`, filling `token` and `bytes` from the tokenizer with the usual U+FFFD byte-fallback correction. Generate returns every top-k candidate the engine produced; derender keeps the ones the coupled endpoint would, with the same helpers: the first `top_logprobs` on chat (all of them for `top_logprobs=-1`) and the first `logprobs + 1` on completions (all of them for `logprobs=-1`). It reads those values from `chat_request` / `completion_request`, so without the original request, or with one that does not ask for logprobs, it returns every candidate. Requests that use `logprob_token_ids` get no output logprobs from generate (their `sampling_params.logprobs` is unset), so derender has none to return for them either. Streaming derender does the same per chunk on both the plain detokenization path and the parser path (see [Streaming state and logprobs](#streaming-state-and-logprobs)). As on `/v1/chat/completions`, a request with `include_reasoning: false` and a parser configured gets no logprobs at all, batch or streaming, since the decoded logprob tokens would reveal the hidden reasoning.
## Streaming cost
@@ -129,7 +129,7 @@ Streaming derender is stateless on the server side: all mutable state lives in t
- `logprob_context_token_ids`: the trailing sampled token IDs (at most 4) from previous chunks, used to seed byte-fallback (U+FFFD) correction so multi-byte characters whose tokens split across chunk boundaries still resolve to real strings
- `logprob_text_offset`: the cumulative emitted text length, so `text_offset` in completion streaming logprobs stays absolute across chunks instead of restarting at 0
When a streamed `GenerateTokensStreamChoice` carries `logprobs`, the integer `token_id`s are decoded per chunk and the resolved logprobs are attached to the corresponding streamed choice. For chat it is as `ChatCompletionLogProbs` and for completions it is converted to the flat `CompletionLogProbs` lists. Chunks without `logprobs` produce choices with `logprobs: null`.
When a streamed `GenerateTokensStreamChoice` carries `logprobs`, the integer `token_id`s are decoded per chunk and the resolved logprobs are attached to the corresponding streamed choice. For chat it is as `ChatCompletionLogProbs` and for completions it is converted to the flat `CompletionLogProbs` lists. Chunks without `logprobs` produce choices with `logprobs: null`. On the parser path the chunk's logprobs are attached to whatever delta the parser emits for that chunk (reasoning, content or a tool call delta, or an empty delta while the parser is still buffering markup), which is what the coupled endpoint does; replay never touches them, since each chunk's `logprobs` cover exactly that chunk's tokens.
## Example
docs/training/async_rl.md
+4
View File
@@ -76,6 +76,10 @@ for span in completion.weight_versions:
The offsets cover the whole output, so with streamed delta outputs, collect the tokens before slicing. `/inference/v1/generate` returns the same list on each choice; see [Tokens In <> Tokens Out API](../serving/online_serving/token_in_token_out.md#weight-versions).
`/v1/completions` and `/v1/chat/completions` return it on each choice when the request sets `return_weight_versions: true`. The offsets count the choice's output tokens, which `return_token_ids` returns as `token_ids`, and a streamed choice gets the list on its last chunk.
An aborted request returns spans over its partial output. If you resubmit that output as part of a new prompt, the new request's spans start again at 0, so shift them by the number of tokens already generated before joining the two lists.
## Example
The [async RLHF example](../../examples/rl/rlhf_async_new_apis.py) demonstrates this pattern with `vllm.AsyncLLMEngine`, NCCL weight transfer, and mid-flight pause/resume with validation.
docs/training/weight_transfer/m2n.md
+18 -13
View File
@@ -10,18 +10,14 @@ The `nccl_m2n` weight transfer engine uses [NCCL M2N](https://github.com/NVIDIA/
## Requirements
This backend depends on the `nccl-extensions` Python package (and its `nccl4py` dependency), which vLLM does not install. It needs NCCL 2.30.5 or newer, and `libnccl_m2n.so` is linked against a specific NCCL build, so vLLM must load the same library:
```bash
export VLLM_NCCL_SO_PATH=/path/to/the/same/libnccl.so
```
This backend depends on the `nccl-extensions` Python package (and its `nccl4py` dependency), which vLLM does not install. Install both against a compatible NCCL 2.30.5 or newer runtime. The backend creates its data-plane communicator through `nccl4py` and passes that communicator directly to `nccl.m2n`.
The backend imports the runtime lazily, so vLLM is unaffected unless you select it.
## How It Works
1. The trainer ranks and all inference workers join **one** NCCL communicator via `StatelessProcessGroup`. Trainer ranks occupy `[0, T)` and workers `[T, T + N)` — m2n meshes are contiguous rank intervals, which is exactly what vLLM's existing `rank_offset` convention produces.
2. At the init handshake the trainer ships the full transfer plan: both meshes — one per side, shared by every parameter — then each parameter's name, dtype, shape and placement. The workers resolve the matching **destination placements** and publish them back over the shared group; the meshes themselves are declared, not negotiated.
1. A `StatelessProcessGroup` exchanges metadata and the NCCL unique ID. The trainer ranks and all inference workers then create **one** `nccl.core.Communicator`: trainer ranks occupy `[0, T)` and workers `[T, T + N)`. The communicator is passed directly to each M2N collective.
2. At the init handshake the trainer ships the full transfer plan: the session-wide destination mesh, then each parameter's name, dtype, shape, source mesh, and source placements. The workers resolve the matching **destination placements** and publish them over the data-plane communicator; the meshes themselves are declared, not negotiated.
3. Each round, both sides issue one `reshard` per parameter in the same order. Every rank in the communicator participates in every reshard, so — as with the NCCL engine — the worker must be inside `update_weights` while the trainer sends. The trainer engine handles that concurrency itself.
### Destination resolution
@@ -83,19 +79,28 @@ engine.send_weights()
`send_weights()` drives the whole round — `start_weight_update`, `update_weights`, `finish_weight_update`. Every rank in the communicator must be inside the same reshard, so `update_weights` has to be in flight while the trainer sends; the engine runs it on a helper thread so the trainer loop does not have to open-code that.
Trainers with a custom producer (a Megatron export, MoE re-fusing) subclass `M2NWeightSource` and supply layouts explicitly.
Trainers with a custom producer (a Megatron export, MoE re-fusing) subclass
`M2NWeightSource` and supply layouts explicitly. A source may yield any
nccl4py-compatible CUDA buffer: a torch tensor, a CUDA Array Interface object,
or a raw device pointer. vLLM keeps buffer owners alive through stream
completion, using allocator stream tracking when the buffer provides it.
Sources must provide buffers on the communicator's current device; nccl4py
resolves their pointers but does not verify device ownership. Raw-pointer
sources are additionally responsible for pointer lifetime and matching their
declared dtype and local shape.
See [`examples/rl/rlhf_m2n.py`](../../../examples/rl/rlhf_m2n.py) for a runnable FSDP → TP example.
## Limitations
- Tensor rank 1..3; 4-D parameters are rejected.
- dtypes: int8/uint8, fp8 e4m3/e5m2, fp16, bf16, int32/uint32, fp32, int64/uint64, fp64. No fp4.
- The trainer's device mesh must cover a contiguous rank interval starting at 0, and be 1-D or 2-D.
- vLLM validates tensor-rank and dtype support while agreeing on the complete
plan; the installed M2N runtime validates them again for each reshard.
- Each parameter's source mesh must cover the same contiguous trainer-rank interval starting at 0, and be 1-D or 2-D.
- `Partial` placements are not supported.
- Models with known model-level checkpoint preprocessing use the full-tensor
fallback for the affected module.
- A reshard cannot be issued from inside a CUDA graph capture (weight updates run outside capture, so this only matters for custom callers).
- The reshard plan is capped per *shard*, not per mesh: at most 16 source shards may feed one destination shard, and one source shard may feed at most 64 destination shards (`MAX_SOURCES` / `MAX_TARGETS` in the m2n build; raising either needs a rebuild).
- A directly-resharded parameter spreads over `tp_size` destination shards, so each one is fed by fewer source shards. A parameter on the fallback path is replicated, i.e. a single destination shard fed by every source shard, so a trainer that **sharded** its parameters is still limited to 16 ranks whenever any parameter falls back — which fused parameters always do. A replicated trainer is unaffected.
- vLLM checks every parameter's shard overlap against conservative 16-source
and 64-target capacities before committing the plan. The installed M2N
runtime remains authoritative and validates each reshard again.
- Each communicator holds a cached staging pool (~256 MiB by default: 4 channels x 64 MiB).
docs/usage/security.md
+2 -2
View File
@@ -423,9 +423,9 @@ Most cache paths default to subdirectories under a single root. Changing `VLLM_C
| Environment Variable | Default | Description |
| --- | --- | --- |
| `VLLM_CACHE_ROOT` | `~/.cache/vllm` | Base cache directory. Respects `XDG_CACHE_HOME` if set. All paths below inherit from this unless explicitly overridden. |
| `VLLM_CACHE_ROOT` | `~/.cache/vllm` | Base directory for vLLM-owned caches. Respects `XDG_CACHE_HOME` if set. |
| *(torch.compile)* | `$VLLM_CACHE_ROOT/torch_compile_cache/` | Compilation cache for AOT-compiled models, Inductor graphs, and Triton kernels. Controlled by `VLLM_DISABLE_COMPILE_CACHE` (set to `1` to disable). |
| `VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR` | `$VLLM_CACHE_ROOT/flashinfer_autotune_cache/<flashinfer-version>/<arch>/<cache-hash>/` | FlashInfer autotune config cache. |
| `VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR` | Managed: `~/.cache/flashinfer/autotune/`; legacy: `$VLLM_CACHE_ROOT/flashinfer_autotune_cache/<flashinfer-version>/<arch>/<cache-hash>/` | Overrides cache placement for both FlashInfer autotune cache APIs. |
| `VLLM_ASSETS_CACHE` | `$VLLM_CACHE_ROOT/assets/` | Downloaded assets (e.g., tokenizer files). |
| `VLLM_XLA_CACHE_PATH` | `$VLLM_CACHE_ROOT/xla_cache/` | XLA/TPU compilation cache. |
| `VLLM_MEDIA_CACHE` | *(disabled)* | Optional cache for downloaded media (images, video, audio). Not enabled unless explicitly set. |
examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
+3 -1
View File
@@ -473,9 +473,11 @@ async def stop_profile():
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=10001)
# Port the prefill/decode servers register on (their proxy_ping_port).
parser.add_argument("--service-discovery-port", type=int, default=36367)
args = parser.parse_args()
t = start_service_discovery("0.0.0.0", 36367)
t = start_service_discovery("0.0.0.0", args.service_discovery_port)
# High-concurrency hardening. Quart's app.run() uses a shallow listen
# backlog (100) and, with app.debug=True, adds per-request overhead that
# starves the single accept loop. Under a burst of ~512 simultaneous client
examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py
+4 -1
View File
@@ -340,7 +340,10 @@ async def _handle_completions(api: str, request: Request):
):
yield chunk
return StreamingResponse(generate_stream(), media_type="application/json")
media_type = (
"text/event-stream" if req_data.get("stream") else "application/json"
)
return StreamingResponse(generate_stream(), media_type=media_type)
except Exception as e:
import sys
examples/rl/rlhf_m2n.py
+25 -10
View File
@@ -17,9 +17,8 @@ Every FSDP rank builds an ``M2NTrainerWeightTransferEngine`` and calls
``send_weights()``; all ranks run every reshard, and only rank 0 drives the
inference side through its ``RayVLLMWeightSyncClient``.
Requires the ``nccl-extensions`` package (NCCL 2.30.5+) and a
``VLLM_NCCL_SO_PATH`` pointing at the same ``libnccl.so`` that
``libnccl_m2n.so`` was linked against.
Requires the ``nccl-extensions`` package (NCCL 2.30.5+). The backend creates
the data-plane communicator with ``nccl4py`` and passes it directly to M2N.
This example was written for 4xH100.
"""
@@ -28,6 +27,7 @@ from __future__ import annotations
import os
import nccl.core as nccl
import pybase64 as base64
import ray
import torch
@@ -38,7 +38,6 @@ from transformers import AutoModelForCausalLM
from vllm import LLM, SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
from vllm.distributed.weight_transfer import (
RayVLLMWeightSyncClient,
WeightTransferTrainerFactory,
@@ -94,7 +93,15 @@ class FSDPTrainWorker:
def ready(self):
return True
def setup_engine(self, llm_handle, nccl_unique_id_b64, world_size, num_workers):
def setup_engine(
self,
llm_handle,
metadata_address,
metadata_port,
nccl_unique_id_b64,
world_size,
num_workers,
):
"""Build the trainer engine on every FSDP rank.
`DTensorModuleSource` reads each parameter's FSDP device mesh and
@@ -103,6 +110,8 @@ class FSDPTrainWorker:
"""
self.engine = WeightTransferTrainerFactory.trainer_init(
init_info=M2NTrainerInitInfo(
master_address=metadata_address,
master_port=metadata_port,
nccl_unique_id_b64=nccl_unique_id_b64,
world_size=world_size,
num_trainer_ranks=FSDP_WORLD_SIZE,
@@ -166,16 +175,22 @@ def main():
# Trainer ranks [0, FSDP_WORLD_SIZE) and inference ranks after them share
# one communicator, so the two meshes are contiguous rank intervals.
nccl = NCCLLibrary()
nccl_unique_id_b64 = base64.b64encode(
bytes(nccl.ncclGetUniqueId().internal)
).decode()
nccl_unique_id_b64 = base64.b64encode(nccl.get_unique_id().as_bytes).decode()
world_size = FSDP_WORLD_SIZE + num_workers
metadata_address = get_ip()
metadata_port = get_open_port()
print("[transfer] Initializing nccl_m2n weight transfer (all FSDP ranks)...")
ray.get(
[
w.setup_engine.remote(llm, nccl_unique_id_b64, world_size, num_workers)
w.setup_engine.remote(
llm,
metadata_address,
metadata_port,
nccl_unique_id_b64,
world_size,
num_workers,
)
for w in fsdp_workers
]
)
rust/src/chat/src/lib.rs
+2 -2
View File
@@ -434,7 +434,7 @@ mod tests {
)
.unwrap_err();
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hf, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_json, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: beam, deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hf, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_json, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
}
#[test]
@@ -445,6 +445,6 @@ mod tests {
)
.unwrap_err();
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hf, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, mimo, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: beam, cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hf, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, mimo, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
}
}
rust/src/chat/src/output/default/mod.rs
+29 -12
View File
@@ -23,7 +23,8 @@ use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEvent
use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory};
use crate::parser::tool::{ToolParser, ToolParserFactory};
use crate::parser::unified::{
HfTemplateError, HfUnifiedParser, ResponseTemplate, UnifiedParserFactory, names,
BeamUnifiedParser, HfTemplateError, HfUnifiedParser, ResponseTemplate, UnifiedParserFactory,
names,
};
use crate::parser::{ParserSelection, ToolStrictLevel};
use crate::request::{ChatRequest, ChatTool};
@@ -206,17 +207,19 @@ impl DefaultChatOutputProcessor {
});
}
// `hf` is built from the model's template, never by its registry constructor.
let parser = if parser_name == names::HF {
let template =
response_template.as_ref().map_err(|error| Error::ParserInitialization {
kind: "unified",
name: parser_name.to_string(),
error: Box::new(error.clone()),
})?;
Box::new(HfUnifiedParser::new(template.clone(), tools, tokenizer))
} else {
factory.create(parser_name, tools, tokenizer)?
// Parsers driven by the model's template are built from it, never by
// their registry constructors.
let template = || {
response_template.clone().map_err(|error| Error::ParserInitialization {
kind: "unified",
name: parser_name.to_string(),
error: Box::new(error),
})
};
let parser: Box<dyn UnifiedParser> = match parser_name {
names::HF => Box::new(HfUnifiedParser::new(template()?, tools, tokenizer)),
names::BEAM => Box::new(BeamUnifiedParser::new(template()?, tools, tokenizer)),
_ => factory.create(parser_name, tools, tokenizer)?,
};
UNIFIED_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using unified parser"));
@@ -589,6 +592,20 @@ mod tests {
.assert_eq(&error(invalid_template()));
}
#[test]
fn explicit_beam_uses_the_model_response_template() {
let beam = ParserSelection::Explicit("beam".to_string());
crate::parser::validate_parser_overrides(&beam, &beam).unwrap();
assert!(keeps_special_tokens("other-model", &loaded_template(), &beam, &beam).unwrap());
let error =
keeps_special_tokens("other-model", &Err(HfTemplateError::Missing), &beam, &beam)
.unwrap_err()
.to_report_string();
expect_test::expect!["failed to initialize unified parser `beam`: the model's tokenizer_config.json provides no response_template"]
.assert_eq(&error);
}
#[test]
fn auto_never_selects_hf() {
let auto = ParserSelection::Auto;
rust/src/chat/src/parser/unified.rs
+6 -2
View File
@@ -6,8 +6,8 @@
use std::sync::LazyLock;
pub use vllm_parser::unified::{
Gemma4UnifiedParser, HfTemplateError, HfUnifiedParser, HyV3UnifiedParser, HyV4UnifiedParser,
InklingUnifiedParser, KimiK3UnifiedParser, ResponseTemplate, UnifiedParser,
BeamUnifiedParser, Gemma4UnifiedParser, HfTemplateError, HfUnifiedParser, HyV3UnifiedParser,
HyV4UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, ResponseTemplate, UnifiedParser,
};
use vllm_tokenizer::DynTokenizer;
@@ -26,6 +26,9 @@ pub mod names {
/// name resolution only: it is built from the model's template, never by its
/// registry constructor.
pub const HF: &str = "hf";
/// Beam checkpoints: [`HF`] with Beam's tool-call grammar, built from the
/// model's template the same way.
pub const BEAM: &str = "beam";
}
/// Constructor signature for one registered unified parser implementation.
@@ -55,6 +58,7 @@ impl UnifiedParserFactory {
factory.register_parser::<KimiK3UnifiedParser>(names::KIMI_K3);
// Opt-in only: no model patterns map to it.
factory.register_parser::<HfUnifiedParser>(names::HF);
factory.register_parser::<BeamUnifiedParser>(names::BEAM);
factory
.register_pattern("gemma-4", names::GEMMA4)
rust/src/engine-core-client/src/client.rs
+2 -5
View File
@@ -657,12 +657,9 @@ impl EngineCoreClient {
);
let result: Result<()> = async {
// The engine that receives it wakes the other ranks if they are idle.
if let Some(coordinator) = self.coordinator.as_ref() {
let snapshot = coordinator.snapshot();
req.current_wave = snapshot.current_wave;
if !snapshot.engines_running {
coordinator.notify_first_request(engine_id.clone())?;
}
req.current_wave = coordinator.snapshot().current_wave;
}
debug!(
rust/src/engine-core-client/src/coordinator/external.rs
+14 -83
View File
@@ -3,34 +3,16 @@
use std::sync::Arc;
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
use serde_tuple::Deserialize_tuple;
use thiserror_ext::AsReport;
use tokio::sync::mpsc;
use tracing::{debug, warn};
use zeromq::prelude::{SocketRecv, SocketSend};
use zeromq::prelude::SocketRecv;
use zeromq::{XSubSocket, ZmqMessage};
use crate::client::imp::ClientInner;
use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState};
use crate::coordinator::handle::CoordinatorState;
use crate::error::{Error, Result, bail_unexpected_coordinator_output};
use crate::protocol::{OpaqueValue, decode_msgpack, encode_msgpack};
/// Frontend-to-coordinator wakeup message sent when the first request arrives
/// while all engines are paused.
///
/// This matches the frontend-side msgpack tuple sent by Python
/// `DPAsyncMPClient._ensure_stats_update_task` to the coordinator front socket.
///
/// Original Python definition:
/// <https://github.com/vllm-project/vllm/blob/694449050f8dac3d9853e97e518b4a43ec52106a/vllm/v1/engine/core_client.py#L1230-L1236>
#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)]
struct CoordinatorWakeupMessage {
/// Engine index that already has the triggering request and should be
/// excluded from the coordinator's `START_DP_WAVE` rebroadcast.
exclude_engine_index: u32,
/// DP wave number observed by the frontend when the request was admitted.
wave: u32,
}
use crate::protocol::{OpaqueValue, decode_msgpack};
/// Coordinator-to-frontend state publish received on the front-side coordinator
/// socket.
@@ -56,56 +38,17 @@ struct CoordinatorStateUpdate {
/// Background half of an external Python-owned coordinator connection.
///
/// This owns the command receiver and one frontend-facing XSUB socket. It
/// mirrors the subset of Python's coordinator protocol needed by the Rust
/// bootstrapped frontend: receive `(counts, wave, running)` publishes, ignore
/// `counts`, and send `(exclude_engine_index, wave)` wakeup messages when the
/// first request arrives while engines are paused.
/// This owns one frontend-facing XSUB socket. It mirrors the subset of Python's
/// coordinator protocol needed by the Rust bootstrapped frontend: receive
/// `(counts, wave, running)` publishes and ignore `counts`.
pub(crate) struct ExternalCoordinatorService {
state: Arc<CoordinatorState>,
command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
socket: XSubSocket,
}
impl ExternalCoordinatorService {
pub(super) fn new(
state: Arc<CoordinatorState>,
command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
socket: XSubSocket,
) -> Self {
Self {
state,
command_rx,
socket,
}
}
/// Apply one frontend-originated command to the external coordinator state
/// machine.
async fn handle_command(&mut self, command: CoordinatorCommand) -> Result<()> {
match command {
CoordinatorCommand::FirstRequest {
target_engine_id,
wave,
} => {
let target_engine_index = target_engine_id.engine_index().ok_or_else(|| {
Error::UnsupportedCoordinatorEngineId {
engine_id: target_engine_id.to_vec(),
}
})?;
debug!(
wave,
exclude_engine_index = target_engine_index,
"notifying external coordinator about first request while engines were paused"
);
let payload = encode_msgpack(&CoordinatorWakeupMessage {
exclude_engine_index: target_engine_index,
wave,
})?;
self.socket.send(ZmqMessage::from(payload)).await?;
}
}
Ok(())
pub(super) fn new(state: Arc<CoordinatorState>, socket: XSubSocket) -> Self {
Self { state, socket }
}
/// Apply one publish received from the xsub socket containing a coordinator
@@ -136,26 +79,14 @@ impl ExternalCoordinatorService {
Ok(())
}
/// Drive the coordinator event loop until either side of the control plane
/// is closed or a fatal error is observed.
/// Drive the coordinator event loop until the task is aborted or a fatal
/// error is observed.
pub(crate) async fn run(mut self, inner: Arc<ClientInner>) {
let result: Result<()> = async {
loop {
tokio::select! {
// Received frontend-originated command from the handle.
command = self.command_rx.recv() => {
let Some(command) = command else {
warn!("external coordinator command channel closed, shutting down service");
return Ok(());
};
self.handle_command(command).await?;
}
// Received publish from the external coordinator socket.
publish = self.socket.recv() => {
let publish = publish.map_err(Error::from)?;
self.handle_publish(publish).await?;
}
}
// Received publish from the external coordinator socket.
let publish = self.socket.recv().await.map_err(Error::from)?;
self.handle_publish(publish).await?;
}
}
.await;
rust/src/engine-core-client/src/coordinator/handle.rs
+10 -81
View File
@@ -4,101 +4,56 @@
use std::sync::Arc;
use parking_lot::Mutex;
use tokio::sync::mpsc;
use zeromq::prelude::Socket;
use zeromq::{XPubSocket, XSubSocket};
use crate::coordinator::external::ExternalCoordinatorService;
use crate::coordinator::inproc::InProcCoordinatorRunner;
use crate::error::{Error, Result, bail_control_closed};
use crate::transport::EngineId;
use crate::error::Result;
/// Snapshot to the coordinator state for request routing and stamping.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct CoordinatorStateSnapshot {
/// The current DP wave, which will be stamped on outgoing requests.
pub current_wave: u32,
/// Whether the engines are currently running or paused, which determines if
/// the frontend must trigger a new wave on the next request.
/// Whether the engines are currently running or paused, as last reported
/// by the engines themselves.
pub engines_running: bool,
}
impl CoordinatorStateSnapshot {
/// Resume the engines for a `FirstRequest` and return the wave to broadcast
/// and the engine to exclude from the wakeup.
///
/// The request may have been stamped with a `request_wave` older than
/// `current_wave` if a `WaveComplete` advanced it after the command was
/// enqueued. Such a request still needs serving, so the current wave is
/// broadcast to every engine (`exclude = None`); the wave is never rewound.
/// A non-stale request excludes the engine that already received it. Mirrors
/// the Python coordinator's front-end path.
pub(crate) fn start_wave_for_first_request(
&mut self,
request_wave: u32,
target_engine_index: u32,
) -> (u32, Option<u32>) {
self.engines_running = true;
let exclude = (request_wave >= self.current_wave).then_some(target_engine_index);
(self.current_wave, exclude)
}
}
/// Shared in-process coordinator state.
pub(crate) type CoordinatorState = Mutex<CoordinatorStateSnapshot>;
/// Commands sent from the frontend request path into the background runner.
#[derive(Debug)]
pub(crate) enum CoordinatorCommand {
/// The first request arrived while all engines were paused.
///
/// The coordinator should broadcast `START_DP_WAVE` with the current wave
/// and the target engine index as the excluded engine.
FirstRequest {
target_engine_id: EngineId,
wave: u32,
},
}
/// Frontend-facing coordinator handle used by `EngineCoreClient::call()`.
///
/// This side stays intentionally small: it can read the latest wave snapshot
/// and enqueue a `FirstRequest` transition when the request path observes the
/// system in the paused state.
/// The frontend only reads the latest wave snapshot: waves are started by the
/// engines, which notify the coordinator when an idle engine receives a request.
#[derive(Clone)]
pub(crate) struct CoordinatorHandle {
state: Arc<CoordinatorState>,
command_tx: mpsc::UnboundedSender<CoordinatorCommand>,
}
impl CoordinatorHandle {
fn new_parts() -> (
Self,
Arc<CoordinatorState>,
mpsc::UnboundedReceiver<CoordinatorCommand>,
) {
fn new_parts() -> (Self, Arc<CoordinatorState>) {
let state = Arc::new(Mutex::new(CoordinatorStateSnapshot {
current_wave: 0,
engines_running: false,
}));
let (command_tx, command_rx) = mpsc::unbounded_channel();
(
Self {
state: state.clone(),
command_tx,
},
state,
command_rx,
)
}
/// Build the paired frontend handle and background runner around one
/// engine-facing coordinator broadcast socket.
pub(crate) fn new_inproc(coordinator_input: XPubSocket) -> (Self, InProcCoordinatorRunner) {
let (handle, state, command_rx) = Self::new_parts();
let (handle, state) = Self::new_parts();
(
handle,
InProcCoordinatorRunner::new(state, command_rx, coordinator_input),
InProcCoordinatorRunner::new(state, coordinator_input),
)
}
@@ -107,41 +62,15 @@ impl CoordinatorHandle {
pub(crate) async fn connect_external(
coordinator_address: &str,
) -> Result<(Self, ExternalCoordinatorService)> {
let (handle, state, command_rx) = Self::new_parts();
let (handle, state) = Self::new_parts();
let mut socket = XSubSocket::new();
socket.connect(coordinator_address).await?;
socket.subscribe("").await?;
Ok((
handle,
ExternalCoordinatorService::new(state, command_rx, socket),
))
Ok((handle, ExternalCoordinatorService::new(state, socket)))
}
/// Snapshot the coordinator state for request routing and stamping.
pub(crate) fn snapshot(&self) -> CoordinatorStateSnapshot {
*self.state.lock()
}
/// Notify the runner that a new request arrived while engines were paused.
///
/// The handle flips `engines_running` optimistically so concurrent request
/// submissions coalesce behind one `START_DP_WAVE` broadcast instead of all
/// trying to trigger the wave independently.
pub(crate) fn notify_first_request(&self, target_engine_id: EngineId) -> Result<()> {
let mut state = self.state.lock();
if state.engines_running {
return Ok(());
}
let command = CoordinatorCommand::FirstRequest {
target_engine_id,
wave: state.current_wave,
};
if self.command_tx.send(command).is_err() {
bail_control_closed!("in-process coordinator command channel already shut down");
}
state.engines_running = true;
Ok(())
}
}
rust/src/engine-core-client/src/coordinator/inproc.rs
+10 -98
View File
@@ -12,7 +12,7 @@ use zeromq::{XPubSocket, ZmqError, ZmqMessage};
use crate::client::imp::ClientInner;
use crate::coordinator::bootstrap::is_engine_subscription;
use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState};
use crate::coordinator::handle::CoordinatorState;
use crate::error::{Error, Result, bail_unexpected_coordinator_output};
use crate::protocol::encode_msgpack;
use crate::protocol::output::{DpControlMessage, DpControlOutput, EngineCoreOutputs};
@@ -30,32 +30,25 @@ use crate::protocol::request::EngineCoreRequestType;
struct StartDpWaveMessage {
/// DP wave number that all engines should start processing.
wave: u32,
/// Engine index that already received the triggering request and so does not
/// need an extra wakeup. `None` wakes every engine (used when the triggering
/// request was for a stale wave).
/// Engine index that started the wave and so does not need an extra wakeup.
/// `None` wakes every engine.
exclude_engine_index: Option<u32>,
}
/// Background half of the in-process coordinator.
///
/// This owns the command receiver and the engine-facing coordinator input
/// socket. It is the single place where wave transitions are serialized and
/// where `START_DP_WAVE` broadcasts are emitted.
/// This owns the engine-facing coordinator input socket. It is the single place
/// where wave transitions are serialized and where `START_DP_WAVE` broadcasts
/// are emitted.
pub(crate) struct InProcCoordinatorRunner {
state: Arc<CoordinatorState>,
command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
coordinator_input: XPubSocket,
}
impl InProcCoordinatorRunner {
pub(super) fn new(
state: Arc<CoordinatorState>,
command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
coordinator_input: XPubSocket,
) -> Self {
pub(super) fn new(state: Arc<CoordinatorState>, coordinator_input: XPubSocket) -> Self {
Self {
state,
command_rx,
coordinator_input,
}
}
@@ -82,34 +75,6 @@ impl InProcCoordinatorRunner {
Ok(())
}
/// Apply one frontend-originated command to the coordinator state machine.
async fn handle_command(&mut self, command: CoordinatorCommand) -> Result<()> {
match command {
CoordinatorCommand::FirstRequest {
target_engine_id,
wave,
} => {
let target_engine_index = target_engine_id.engine_index().ok_or_else(|| {
Error::UnsupportedCoordinatorEngineId {
engine_id: target_engine_id.to_vec(),
}
})?;
let (current_wave, exclude) = {
let mut state = self.state.lock();
state.start_wave_for_first_request(wave, target_engine_index)
};
debug!(
current_wave,
request_wave = wave,
?exclude,
"starting DP wave after first request while engines were paused"
);
self.broadcast_start_wave(current_wave, exclude).await?;
}
}
Ok(())
}
/// Apply one engine-originated control output to the coordinator state
/// machine.
async fn handle_outputs(&mut self, outputs: EngineCoreOutputs) -> Result<()> {
@@ -160,7 +125,7 @@ impl InProcCoordinatorRunner {
debug!(
wave,
exclude_engine_index = engine_index,
"starting DP wave after stale-wave notification from engine"
"starting DP wave on notification from engine"
);
self.broadcast_start_wave(wave, Some(engine_index)).await?;
}
@@ -175,8 +140,8 @@ impl InProcCoordinatorRunner {
Ok(())
}
/// Drive the coordinator event loop until either side of the control plane
/// is closed or a fatal error is observed.
/// Drive the coordinator event loop until the engine control output is
/// closed or a fatal error is observed.
///
/// Any fatal error closes the main client registries so request streams and
/// future calls observe a stable shutdown cause.
@@ -216,14 +181,6 @@ impl InProcCoordinatorRunner {
Err(error) => return Err(error.into()),
}
}
// Received frontend-originated command from the handle.
command = self.command_rx.recv() => {
let Some(command) = command else {
warn!("coordinator command channel closed, shutting down coordinator runner");
return Ok(());
};
self.handle_command(command).await?;
}
// Received engine-originated control output from the coordinator socket.
outputs = output_rx.recv() => {
let Some(outputs) = outputs else {
@@ -242,48 +199,3 @@ impl InProcCoordinatorRunner {
inner.close_registries(Arc::new(error));
}
}
#[cfg(test)]
mod tests {
use crate::coordinator::handle::CoordinatorStateSnapshot;
/// A `FirstRequest` for the current wave starts that wave and excludes the
/// engine that already received the triggering request.
#[test]
fn first_request_for_current_wave_excludes_target() {
let mut state = CoordinatorStateSnapshot {
current_wave: 3,
engines_running: false,
};
let (wave, exclude) = state.start_wave_for_first_request(3, 2);
assert_eq!(wave, 3);
assert_eq!(exclude, Some(2));
assert!(state.engines_running);
assert_eq!(state.current_wave, 3);
}
/// A `FirstRequest` whose wave was superseded by a racing `WaveComplete`
/// (`request_wave < current_wave`) must still start the request's wave: it
/// broadcasts the current wave and wakes every engine (`exclude = None`)
/// rather than rewinding the wave or dropping the request.
#[test]
fn stale_first_request_starts_current_wave_for_all_engines() {
let mut state = CoordinatorStateSnapshot {
current_wave: 4,
engines_running: false,
};
// Request stamped with wave 3 while the coordinator already advanced to 4.
let (wave, exclude) = state.start_wave_for_first_request(3, 2);
assert_eq!(
wave, 4,
"must broadcast the current wave, not the stale one"
);
assert_eq!(exclude, None, "a stale request must wake every engine");
assert!(state.engines_running);
assert_eq!(state.current_wave, 4, "wave must not be rewound");
}
}
rust/src/engine-core-client/src/error.rs
-4
View File
@@ -55,8 +55,6 @@ pub enum Error {
UnexpectedCoordinatorOutput { message: String },
#[error("unexpected output on main dispatcher path: {message}")]
UnexpectedDispatcherOutput { message: String },
#[error("coordinator requires a Python-compatible two-byte engine id, got {engine_id:?}")]
UnsupportedCoordinatorEngineId { engine_id: Vec<u8> },
#[error("unsupported auxiliary frame(s): expected 1 frame, got {frame_count}")]
UnsupportedAuxFrames { frame_count: usize },
#[error("external coordinator mode is not implemented yet")]
@@ -66,8 +64,6 @@ pub enum Error {
context: &'static str,
field: &'static str,
},
#[error("engine control channel closed unexpectedly: {message}")]
ControlClosed { message: String },
#[error("request `{request_id}` is already in flight")]
DuplicateRequestId { request_id: String },
#[error(
rust/src/engine-core-client/src/tests/client.rs
+28 -103
View File
@@ -267,6 +267,19 @@ async fn recv_start_dp_wave(sub: &mut SubSocket) -> (u32, Option<u32>) {
rmp_serde::from_slice(&frames[1]).expect("decode START_DP_WAVE payload")
}
async fn send_start_wave(push: &mut PushSocket, engine_index: u32, wave: u32) {
send_outputs(
push,
DpControlOutput {
engine_index,
timestamp: 0.0,
control: DpControlMessage::StartWave(wave),
}
.into(),
)
.await;
}
async fn connect_client_with_ipc(
config: EngineCoreClientConfig,
ipc: &IpcNamespace,
@@ -413,12 +426,6 @@ async fn recv_xpub_subscription(xpub: &mut XPubSocket) {
assert_eq!(frames[0].as_ref(), b"\x01");
}
async fn recv_external_coordinator_wakeup(xpub: &mut XPubSocket) -> (u32, u32) {
let frames = recv_xpub_message(xpub).await;
assert_eq!(frames.len(), 1);
rmp_serde::from_slice(&frames[0]).expect("decode external coordinator wakeup")
}
async fn send_external_coordinator_publish<T: serde::Serialize>(
xpub: &mut XPubSocket,
payload: &T,
@@ -617,6 +624,14 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
engine.coordinator.take().expect("coordinator sockets should be present");
let data_socket = engine.data_sockets.first_mut().expect("data socket");
// An idle engine handed a request starts the wave itself.
let add = recv_engine_message(&mut data_socket.dealer).await;
assert_eq!(add[0].as_ref(), &[0x00]);
let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
assert_eq!(request.request_id, "req-1");
assert_eq!(request.current_wave, 0);
send_start_wave(&mut coordinator.output_push, 0, 0).await;
let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
assert_eq!((wave, exclude_engine), (0, Some(0)));
@@ -629,12 +644,6 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
let replay = recv_start_dp_wave(&mut coordinator.input_sub).await;
assert_eq!(replay, (0, None));
let add = recv_engine_message(&mut data_socket.dealer).await;
assert_eq!(add[0].as_ref(), &[0x00]);
let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
assert_eq!(request.request_id, "req-1");
assert_eq!(request.current_wave, 0);
assert!(
timeout(
Duration::from_millis(200),
@@ -670,14 +679,15 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
)
.await;
let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
assert_eq!((wave, exclude_engine), (1, Some(0)));
let add = recv_engine_message(&mut data_socket.dealer).await;
assert_eq!(add[0].as_ref(), &[0x00]);
let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
assert_eq!(request.request_id, "req-3");
assert_eq!(request.current_wave, 1);
send_start_wave(&mut coordinator.output_push, 0, 1).await;
let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
assert_eq!((wave, exclude_engine), (1, Some(0)));
send_outputs(
&mut data_socket.push,
@@ -896,9 +906,6 @@ async fn coordinator_accepts_stats_only_outputs() {
engine.coordinator.take().expect("coordinator sockets should be present");
let data_socket = engine.data_sockets.first_mut().expect("data socket");
let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
assert_eq!((wave, exclude_engine), (0, Some(0)));
send_outputs(
&mut coordinator.output_push,
RequestBatchOutputs {
@@ -3248,7 +3255,7 @@ async fn bootstrapped_external_coordinator_connects_and_subscribes() {
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends_one_wakeup() {
async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends_no_wakeup() {
init_tracing();
let ipc = IpcNamespace::new().unwrap();
let input_address = ipc.input_endpoint();
@@ -3289,18 +3296,11 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends
let mut stream = client.call(sample_request()).await.unwrap().into_outputs();
let wakeup = timeout(
Duration::from_secs(1),
recv_external_coordinator_wakeup(&mut stats_socket),
)
.await
.unwrap();
assert_eq!(wakeup, (0, 7));
// The engine that receives the request starts the wave, not the frontend.
assert!(
timeout(
Duration::from_millis(200),
recv_external_coordinator_wakeup(&mut stats_socket)
recv_xpub_message(&mut stats_socket)
)
.await
.is_err()
@@ -3333,78 +3333,3 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends
client.shutdown().await.unwrap();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn bootstrapped_external_coordinator_running_state_suppresses_wakeup() {
init_tracing();
let ipc = IpcNamespace::new().unwrap();
let input_address = ipc.input_endpoint();
let output_address = ipc.output_endpoint();
let coordinator_address = ipc.endpoint("stats.sock");
let mut stats_socket = XPubSocket::new();
stats_socket.bind(&coordinator_address).await.unwrap();
let client_task = tokio::spawn({
let input_address = input_address.clone();
let output_address = output_address.clone();
let coordinator_address = coordinator_address.clone();
async move {
EngineCoreClient::connect(bootstrapped_test_config(
input_address,
output_address,
1,
Duration::from_secs(2),
0,
Some(CoordinatorMode::External {
address: coordinator_address,
}),
))
.await
.unwrap()
}
});
let (mut dealer, mut push) =
setup_bootstrapped_mock_engine(input_address, output_address, &[0x00, 0x00]).await;
let client = client_task.await.unwrap();
recv_xpub_subscription(&mut stats_socket).await;
send_external_coordinator_publish(&mut stats_socket, &(Value::Nil, 5_u32, true)).await;
tokio::time::sleep(Duration::from_millis(50)).await;
let mut stream = client.call(sample_request()).await.unwrap().into_outputs();
assert!(
timeout(
Duration::from_millis(200),
recv_external_coordinator_wakeup(&mut stats_socket)
)
.await
.is_err()
);
let add = recv_engine_message(&mut dealer).await;
let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
assert_eq!(request.current_wave, 5);
send_outputs(
&mut push,
RequestBatchOutputs {
outputs: vec![request_output(
"req-1",
vec![],
Some(EngineCoreFinishReason::Length),
)],
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
..Default::default()
}
.into(),
)
.await;
let final_output = timeout(Duration::from_secs(1), stream.next()).await.unwrap();
assert!(final_output.is_some());
client.shutdown().await.unwrap();
}
rust/src/parser/src/output_grammar.rs
+3 -2
View File
@@ -207,7 +207,8 @@ pub(crate) fn full_format_from_builder_for_test(
)
}
fn tool_params(tools: &[Tool], strict_level: ToolStrictLevel) -> Vec<ToolParam> {
/// Tools as builder inputs, strict when the tool or the server floor asks for it.
pub(crate) fn tool_params(tools: &[Tool], strict_level: ToolStrictLevel) -> Vec<ToolParam> {
tools
.iter()
.map(|tool| {
@@ -240,7 +241,7 @@ pub(crate) fn normalize_tool_choice(
}
/// Whether the request asks for a tool grammar at all.
fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool {
pub(crate) fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool {
if ctx.tools.is_empty() {
return false;
}
rust/src/parser/src/output_grammar/arguments.rs
+77 -20
View File
@@ -26,7 +26,7 @@
//! - `allOf` with more than one schema accepts any value.
use serde_json::{Value, json};
use xgrammar_structural_tag::format::{Format, JsonSchemaFormat};
use xgrammar_structural_tag::format::{EndBoundary, Format, JsonSchemaFormat, TokenFormat};
pub use crate::schema::JsonType;
use crate::schema::{self, OptionSource, SchemaOption, SchemaRoot};
@@ -60,7 +60,13 @@ impl ParameterKey<'_> {
/// `prefix KEY suffix content end`: a tag for a declared key, or a
/// sequence matching [`Self::FREE_PATTERN`] for a free key.
pub fn tag(&self, prefix: &str, suffix: &str, content: Format, end: &str) -> Format {
pub fn tag(
&self,
prefix: &str,
suffix: &str,
content: Format,
end: impl Into<EndBoundary>,
) -> Format {
match self {
Self::Declared(key) => Format::tag(format!("{prefix}{key}{suffix}"), content, end),
Self::Free => Format::sequence(vec![
@@ -68,7 +74,13 @@ impl ParameterKey<'_> {
Format::regex(Self::FREE_PATTERN),
Format::const_string(suffix),
content,
Format::const_string(end),
match end.into() {
EndBoundary::Text(end) => Format::const_string(end),
EndBoundary::Texts(ends) => {
one_of(ends.into_iter().map(Format::const_string).collect())
}
EndBoundary::Token(end) => Format::Token(TokenFormat { token: end.token }),
},
]),
}
}
@@ -183,6 +195,36 @@ impl ArgumentContext<'_> {
/// Build the grammar of one call's arguments under `schema`, the tool's
/// parameters schema, rendering each parameter with `syntax`.
pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &ArgumentOptions) -> Format {
parameter_list(
options,
syntax.separator(),
parameters(schema, syntax, options),
)
}
/// The parameters of one call, each rendered by an [`ArgumentSyntax`], for a
/// syntax that lists them itself rather than through [`arguments`].
pub struct Parameters {
/// The declared parameters, in declared order.
pub declared: Vec<Parameter>,
/// Any one additional parameter, if the schema allows them.
pub additional: Option<Format>,
}
/// One declared parameter.
pub struct Parameter {
pub key: String,
pub format: Format,
pub required: bool,
}
/// The parameters of one call under `schema`, the tool's parameters schema,
/// rendering each parameter with `syntax`.
pub fn parameters(
schema: &Value,
syntax: &dyn ArgumentSyntax,
options: &ArgumentOptions,
) -> Parameters {
let cx = ArgumentContext {
options,
root: SchemaRoot::new(schema),
@@ -190,8 +232,7 @@ pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &Argument
let parameter =
|key: ParameterKey<'_>, schema: &Value| syntax.parameter(key, &cx.value_options(schema));
let schema = cx.root.resolved();
let (properties, additional) = match schema {
match cx.root.resolved() {
Value::Object(schema) => {
let additional = match schema.get("additionalProperties") {
Some(Value::Bool(false)) => None,
@@ -206,22 +247,33 @@ pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &Argument
.and_then(Value::as_array)
.map(|required| required.iter().filter_map(Value::as_str).collect::<Vec<_>>())
.unwrap_or_default();
let properties = schema
let declared = schema
.get("properties")
.and_then(Value::as_object)
.into_iter()
.flatten()
.filter_map(|(key, schema)| {
let parameter = parameter(ParameterKey::Declared(key), schema)?;
Some((parameter, required.contains(&key.as_str())))
Some(Parameter {
key: key.clone(),
format: parameter(ParameterKey::Declared(key), schema)?,
required: required.contains(&key.as_str()),
})
})
.collect();
(properties, additional)
Parameters {
declared,
additional,
}
}
Value::Bool(false) => (vec![], None),
_ => (vec![], parameter(ParameterKey::Free, &ANY_SCHEMA)),
};
parameter_list(options, syntax.separator(), properties, additional)
Value::Bool(false) => Parameters {
declared: vec![],
additional: None,
},
_ => Parameters {
declared: vec![],
additional: parameter(ParameterKey::Free, &ANY_SCHEMA),
},
}
}
/// Parameters in declared order, required ones mandatory, followed by any
@@ -230,8 +282,10 @@ pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &Argument
fn parameter_list(
options: &ArgumentOptions,
separator: Option<Format>,
properties: Vec<(Format, bool)>,
additional: Option<Format>,
Parameters {
declared,
additional,
}: Parameters,
) -> Format {
let separated = |parameter: Format| match &separator {
Some(separator) => Format::sequence(vec![parameter, separator.clone()]),
@@ -239,10 +293,10 @@ fn parameter_list(
};
let mut elements = Vec::from_iter(separator.clone());
if options.any_order {
let any_required = properties.iter().any(|(_, required)| *required);
let parameters = properties
let any_required = declared.iter().any(|parameter| parameter.required);
let parameters = declared
.into_iter()
.map(|(parameter, _)| parameter)
.map(|parameter| parameter.format)
.chain(additional)
.collect::<Vec<_>>();
if !parameters.is_empty() {
@@ -254,8 +308,11 @@ fn parameter_list(
});
}
} else {
for (parameter, required) in properties {
let parameter = separated(parameter);
for Parameter {
format, required, ..
} in declared
{
let parameter = separated(format);
elements.push(if required {
parameter
} else {
rust/src/parser/src/unified/beam.rs
+543
View File
@@ -0,0 +1,543 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//! Unified parser for Beam checkpoints.
//!
//! Beam output is parsed by [`HfUnifiedParser`] with the checkpoint's
//! `response_template`; this parser adds Beam's output grammar.
//! A Beam message after `<|start|>assistant` is:
//!
//! ```text
//! <|channel|>analysis<|message|>REASONING<|end|><|start|>assistant
//! ( <|channel|>final<|message|>ANSWER<|return|>
//! | CALL (<|end|><|start|>assistant CALL)* <|call|> )
//! ```
//!
//! where reasoning may end at any closer of the template's reasoning field,
//! and `CALL` is
//! ` to=NAME<|channel|>analysis <|constrain|>xml<|message|>` followed by
//! `<rfl:parameter name="KEY">\nVALUE\n</rfl:parameter>` parameters separated
//! by newlines. Where the template looks ahead for the next parameter, a value
//! may also end on a newline right before it.
//!
//! The grammar enforces the schema of every declared tool's arguments, in
//! declared or alphabetical order, since Beam drifts from schemas it does not
//! enforce. When the model may answer instead, calls may also name any other
//! tool with any arguments, which the parser passes on so a harness can tell
//! the model the tool does not exist. `ANSWER` follows the request's structured
//! output, if any.
//!
//! The chat template rejects a history turn without reasoning, so a message the
//! model opens holds non-empty reasoning: an empty one would come back as null
//! reasoning and fail the next request.
//!
//! Unless `VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT` turns it off, the grammar
//! applies to every message, including one the request does not constrain.
use std::collections::BTreeMap;
use std::sync::{Arc, LazyLock};
use serde_json::Value;
use vllm_tokenizer::{DecodedText, DynTokenizer};
use xgrammar_structural_tag::ToolChoice;
use xgrammar_structural_tag::format::{EndBoundary, Format};
use xgrammar_structural_tag::tool::{
BuilderToolChoice, NormalizedToolChoice, ToolChoiceValue, normalize_tool_choice,
};
use super::{
HfPosition, HfUnifiedParser, ResponseTemplate, Result, UnifiedParser, UnifiedParserError,
UnifiedParserOutput,
};
use crate::output_grammar::arguments::{
self, ArgumentOptions, ArgumentSyntax, JsonType, ParameterKey, Parameters, ValueOption, one_of,
};
use crate::output_grammar::{
self, BuiltOutputGrammar, OutputGrammarContext, tool_grammar_applies, tool_params,
};
use crate::tool::Tool;
const ANALYSIS_OPEN: &str = "<|channel|>analysis<|message|>";
const FINAL_OPEN: &str = "<|channel|>final<|message|>";
const END: &str = "<|end|>";
/// Closes one message and opens the next assistant message.
const NEXT_MESSAGE: &str = "<|end|><|start|>assistant";
const RETURN: &str = "<|return|>";
const CALL: &str = "<|call|>";
/// Follows ` to=NAME` in a call.
const CALL_HEADER_END: &str = "<|channel|>analysis <|constrain|>xml<|message|>";
const PARAMETER_OPEN: &str = "<rfl:parameter name=\"";
const PARAMETER_NAME_CLOSE: &str = "\">\n";
const PARAMETER_CLOSE: &str = "\n</rfl:parameter>";
/// Ends a value where the template looks ahead for the next parameter.
const NEXT_PARAMETER: &str = "\n<rfl:parameter name=\"";
/// Ends of a lenient value that another parameter follows: the closer and
/// separator before the next opener, or the opener right after the value.
const NEXT_PARAMETER_ENDS: [&str; 3] = [
"\n</rfl:parameter><rfl:parameter name=\"",
"\n</rfl:parameter>\n<rfl:parameter name=\"",
NEXT_PARAMETER,
];
/// Text that ends the call in the parser, even inside a JSON value.
const CALL_TERMINATORS: [&str; 2] = [END, CALL];
/// Whether every Beam message follows the grammar, even when the request
/// constrains neither tool calls nor the answer. Read from
/// `VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT` like the Python frontend: on unless
/// set to something other than `true` or `1`.
static ALWAYS_CONSTRAIN_OUTPUT: LazyLock<bool> = LazyLock::new(|| {
std::env::var("VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT").map_or(true, |value| {
matches!(value.to_lowercase().as_str(), "true" | "1")
})
});
/// Beam parser: [`HfUnifiedParser`] with Beam's output grammar.
pub struct BeamUnifiedParser {
inner: HfUnifiedParser,
/// Whether a value may end where the next parameter opens.
lenient: bool,
/// What ends the reasoning and opens the next message.
reasoning_ends: Vec<String>,
}
impl BeamUnifiedParser {
/// Create a parser for one request from the checkpoint's `response_template`.
pub fn new(template: Arc<ResponseTemplate>, tools: &[Tool], tokenizer: DynTokenizer) -> Self {
Self {
lenient: template.argument_lookahead() == Some(PARAMETER_OPEN),
reasoning_ends: match template.reasoning_closers() {
[] => vec![NEXT_MESSAGE.to_string()],
closers => closers.to_vec(),
},
inner: HfUnifiedParser::new(template, tools, tokenizer),
}
}
}
impl UnifiedParser for BeamUnifiedParser {
fn create(_tools: &[Tool], _tokenizer: DynTokenizer) -> Result<Box<dyn UnifiedParser>> {
Err(UnifiedParserError::NoNamedConstructor {
parser: "beam",
built_from: "the checkpoint's response_template",
})
}
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
self.inner.initialize(prompt_token_ids)
}
fn preserve_special_tokens(&self) -> bool {
self.inner.preserve_special_tokens()
}
fn build_output_grammar(
&self,
ctx: &OutputGrammarContext<'_>,
) -> output_grammar::Result<Option<BuiltOutputGrammar>> {
build_output_grammar(
self.inner.position(),
ctx,
*ALWAYS_CONSTRAIN_OUTPUT,
self.lenient,
&self.reasoning_ends,
)
}
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
self.inner.tool_call_id(tool_index)
}
fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> {
self.inner.parse_into(delta, output)
}
fn finish(&mut self) -> Result<UnifiedParserOutput> {
self.inner.finish()
}
fn reset(&mut self) -> String {
self.inner.reset()
}
}
/// The grammar of a Beam message from the first generated token, or `None`
/// when the parser is neither at the start of the message nor inside its
/// reasoning, or when the request constrains nothing and `always_constrain` is
/// off. Values may end where the next parameter opens if `lenient`, and
/// reasoning ends at any of `reasoning_ends`.
fn build_output_grammar(
position: HfPosition,
ctx: &OutputGrammarContext<'_>,
always_constrain: bool,
lenient: bool,
reasoning_ends: &[String],
) -> output_grammar::Result<Option<BuiltOutputGrammar>> {
let forced = !ctx.tools.is_empty()
&& matches!(
ctx.tool_choice,
ToolChoice::Value(ToolChoiceValue::Required) | ToolChoice::NamedFunction(_)
);
let answer_format = ctx.answer.filter(|_| !forced);
// The prompt may still list the tools, so only the grammar keeps the model
// from calling them.
let tools_forbidden = !ctx.tools.is_empty()
&& matches!(ctx.tool_choice, ToolChoice::Value(ToolChoiceValue::None));
if !(always_constrain
|| tool_grammar_applies(ctx)
|| tools_forbidden
|| answer_format.is_some())
{
return Ok(None);
}
let text = Format::any_text();
let reasoning_end = EndBoundary::from(reasoning_ends.to_vec());
let reasoning = match position {
HfPosition::ResponseStart => Format::tag(
ANALYSIS_OPEN,
// At least one character, which cannot start Beam's closers.
Format::sequence(vec![Format::regex("[^<]"), text]),
reasoning_end,
),
HfPosition::Reasoning => Format::tag("", text, reasoning_end),
// Inside an answer the prompt opened, only the always-on grammar
// applies, finishing a free answer. A tool choice that forces a call
// cannot be met there, and the request's own constraint applies as is.
HfPosition::Content => {
return Ok(
(always_constrain && !forced && answer_format.is_none()).then(|| {
BuiltOutputGrammar::from_token_zero(Format::sequence(answer_rest(None)))
}),
);
}
HfPosition::Other => return Ok(None),
};
let NormalizedToolChoice {
function_tools,
choice,
..
} = normalize_tool_choice(
&tool_params(ctx.tools, ctx.tool_strict_level),
ctx.tool_choice.clone(),
)?;
if function_tools.is_empty() {
return Ok(match choice {
// `none`, or no tools to call: the model can only answer.
BuilderToolChoice::Auto => {
Some(BuiltOutputGrammar::from_token_zero(Format::sequence(vec![
reasoning,
answer(answer_format),
])))
}
BuilderToolChoice::Required | BuilderToolChoice::Forced => None,
});
}
let mut calls = function_tools
.iter()
.map(|tool| {
let header =
Format::const_string(format!(" to={}{CALL_HEADER_END}", tool.function.name));
let schema = tool.function.parameters.clone().unwrap_or(Value::Bool(true));
Format::sequence(vec![header, arguments(&schema, lenient)])
})
.collect::<Vec<_>>();
if matches!(choice, BuilderToolChoice::Auto) {
let names = ctx.tools.iter().map(|tool| tool.name.as_str()).collect::<Vec<_>>();
let header = Format::sequence(vec![
Format::const_string(" to="),
Format::regex(other_names(&names)),
Format::const_string(CALL_HEADER_END),
]);
calls.push(Format::sequence(vec![
header,
arguments(&Value::Bool(true), lenient),
]));
}
let call = one_of(calls);
let mut calls = vec![call.clone()];
if ctx.parallel_tool_calls && !choice.is_forced() {
calls.push(Format::star(Format::sequence(vec![
Format::const_string(NEXT_MESSAGE),
call,
])));
}
// The last call ends with the `<|call|>` stop token.
calls.push(Format::optional(Format::const_string(CALL)));
let calls = Format::sequence(calls);
let output = match choice {
BuilderToolChoice::Auto => Format::or(vec![calls, answer(answer_format)]),
BuilderToolChoice::Required | BuilderToolChoice::Forced => calls,
};
Ok(Some(BuiltOutputGrammar::from_token_zero(Format::sequence(
vec![reasoning, output],
))))
}
/// A final answer, ending with the `<|return|>` stop token.
fn answer(format: Option<&Format>) -> Format {
Format::sequence([vec![Format::const_string(FINAL_OPEN)], answer_rest(format)].concat())
}
/// The rest of an open final answer, holding `format` or else free text.
fn answer_rest(format: Option<&Format>) -> Vec<Format> {
vec![
// Without `<|end|>`, no unconstrained message can follow the answer.
format.map_or_else(|| Format::any_text_excluding(&[END, RETURN]), Format::clone),
Format::optional(Format::const_string(RETURN)),
]
}
/// The arguments under `schema`: its declared parameters in declared or
/// alphabetical order, required ones mandatory, then any additional ones. If
/// `lenient`, each value but the last may end by opening the next parameter.
fn arguments(schema: &Value, lenient: bool) -> Format {
let options = ArgumentOptions {
excludes: CALL_TERMINATORS.map(str::to_string).to_vec(),
..ArgumentOptions::default()
};
let render = |syntax: BeamArguments| arguments::parameters(schema, &syntax, &options);
if lenient {
lenient_arguments(render(BeamArguments::Joined), render(BeamArguments::Last))
} else {
closed_arguments(render(BeamArguments::Closed))
}
}
/// Arguments whose values all end with a closer.
fn closed_arguments(
Parameters {
declared,
additional,
}: Parameters,
) -> Format {
let additional = additional.map(Format::star);
let lists = orders(&declared, |parameter| &parameter.key)
.into_iter()
.map(|order| {
let listed = order
.into_iter()
.map(|parameter| listed(parameter.format.clone(), parameter.required));
listed.chain(additional.clone()).collect::<Vec<_>>()
})
.filter(|list| !list.is_empty())
.map(Format::sequence)
.collect::<Vec<_>>();
if lists.is_empty() {
Format::const_string("")
} else {
one_of(lists)
}
}
/// A declared parameter as one another parameter follows, `joined`, and as the
/// last one.
struct LenientParameter {
key: String,
joined: Format,
last: Format,
required: bool,
}
impl LenientParameter {
fn joined(&self) -> Format {
listed(self.joined.clone(), self.required)
}
}
/// Arguments whose values end by opening the next parameter, but for the last
/// one, from the parameters as [`BeamArguments::Joined`] and as
/// [`BeamArguments::Last`].
fn lenient_arguments(joined: Parameters, last: Parameters) -> Format {
// Both forms share their terminators, so they keep the same parameters.
let declared = joined
.declared
.into_iter()
.zip(last.declared)
.map(|(joined, last)| LenientParameter {
key: joined.key,
joined: joined.format,
last: last.format,
required: joined.required,
})
.collect::<Vec<_>>();
let additional = joined.additional.zip(last.additional);
let mut lists = vec![];
for order in orders(&declared, |parameter| &parameter.key) {
// The last parameter is a declared one that follows every required
// one, or an additional one.
if !order.is_empty() {
let first = order.iter().rposition(|parameter| parameter.required).unwrap_or(0);
let mut list =
order[..first].iter().map(|parameter| parameter.joined()).collect::<Vec<_>>();
list.push(last_present(&order[first..]));
lists.push(Format::sequence(list));
}
if let Some((joined, last)) = &additional {
let mut list = order.iter().map(|parameter| parameter.joined()).collect::<Vec<_>>();
list.extend([Format::star(joined.clone()), last.clone()]);
lists.push(Format::sequence(list));
}
}
if lists.is_empty() {
return Format::const_string("");
}
// Each parameter starts past the opener, as the value before writes it.
let arguments = Format::sequence(vec![
separator(),
Format::const_string(PARAMETER_OPEN),
one_of(lists),
]);
if declared.iter().any(|parameter| parameter.required) {
arguments
} else {
Format::optional(arguments)
}
}
/// `parameters`, at least one, in order and the last present one in its last
/// form. Only the first may be required.
fn last_present(parameters: &[&LenientParameter]) -> Format {
if let [only] = parameters {
return only.last.clone();
}
// The last one is in the first half, or after all of it. Halving keeps the
// grammar O(n log n) in size and O(log n) deep.
let (head, tail) = parameters.split_at(parameters.len() / 2);
let mut after_head = head.iter().map(|parameter| parameter.joined()).collect::<Vec<_>>();
after_head.push(last_present(tail));
Format::or(vec![last_present(head), Format::sequence(after_head)])
}
/// `parameters` in declared order, and in alphabetical order by `key` if that
/// differs, since Beam often writes them sorted.
fn orders<P>(parameters: &[P], key: impl Fn(&P) -> &str) -> Vec<Vec<&P>> {
let declared = parameters.iter().collect::<Vec<_>>();
let mut sorted = declared.clone();
sorted.sort_by_key(|parameter| key(parameter));
if sorted
.iter()
.map(|parameter| key(parameter))
.eq(declared.iter().map(|parameter| key(parameter)))
{
vec![declared]
} else {
vec![declared, sorted]
}
}
/// `parameter` as listed: optional unless `required`.
fn listed(parameter: Format, required: bool) -> Format {
if required {
parameter
} else {
Format::optional(parameter)
}
}
fn separator() -> Format {
Format::optional(Format::const_string("\n"))
}
/// A regex for the tool names the call header can spell (no whitespace or
/// `<`) besides `names`.
fn other_names(names: &[&str]) -> String {
#[derive(Default)]
struct Node {
children: BTreeMap<char, Node>,
end: bool,
}
fn escape(char: char, meta: &str) -> String {
if meta.contains(char) {
format!("\\{char}")
} else {
char.to_string()
}
}
/// A name character other than `chars`.
fn other<'a>(chars: impl Iterator<Item = &'a char>) -> String {
let chars = chars.map(|&char| escape(char, r"\]^-[")).collect::<String>();
format!(r"[^\s<{chars}]")
}
/// The rest of a name past `node` that ends on none of `names`.
fn rest(node: &Node, first: bool) -> String {
let any = other([].iter());
let mut branches = vec![format!("{}{any}*", other(node.children.keys()))];
branches.extend(node.children.iter().map(|(&char, child)| {
format!("{}{}", escape(char, r".^$*+?{}[]\|()"), rest(child, false))
}));
let optional = if first || node.end { "" } else { "?" };
format!("(?:{}){optional}", branches.join("|"))
}
let mut root = Node::default();
for name in names {
let mut node = &mut root;
for char in name.chars() {
node = node.children.entry(char).or_default();
}
node.end = true;
}
rest(&root, true)
}
/// How one parameter is written: `<rfl:parameter name="KEY">\nVALUE` followed
/// by its end.
///
/// The parser reads every value as text and converts it by the tool schema, so
/// string options are raw text and every other option is JSON.
enum BeamArguments {
/// The whole parameter, ending with `\n</rfl:parameter>`, after an optional
/// newline.
Closed,
/// The parameter past its opener, where the template looks ahead for the
/// next one: it ends by opening the next parameter, also right after the
/// value's newline, so raw text holds no such opener.
Joined,
/// The last parameter past its opener under such a template, ending with
/// `\n</rfl:parameter>`.
Last,
}
impl ArgumentSyntax for BeamArguments {
fn separator(&self) -> Option<Format> {
None
}
fn parameter(&self, key: ParameterKey<'_>, options: &[ValueOption<'_>]) -> Option<Format> {
let terminators: &[&str] = match self {
Self::Closed => &["</rfl:parameter>"],
Self::Joined | Self::Last => &["</rfl:parameter>", NEXT_PARAMETER],
};
let values = options
.iter()
.filter_map(|option| match option.ty {
JsonType::String => option.raw_string(terminators),
_ => Some(option.json()),
})
.collect::<Vec<_>>();
(!values.is_empty()).then(|| {
let value = one_of(values);
match self {
Self::Closed => Format::sequence(vec![
separator(),
key.tag(PARAMETER_OPEN, PARAMETER_NAME_CLOSE, value, PARAMETER_CLOSE),
]),
Self::Joined => key.tag(
"",
PARAMETER_NAME_CLOSE,
value,
NEXT_PARAMETER_ENDS.to_vec(),
),
Self::Last => key.tag("", PARAMETER_NAME_CLOSE, value, PARAMETER_CLOSE),
}
})
}
}
#[cfg(test)]
mod tests;
rust/src/parser/src/unified/beam/tests.rs
+897
View File
@@ -0,0 +1,897 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::sync::Arc;
use serde_json::{Value, json};
use vllm_tokenizer::test_utils::TestTokenizer;
use vllm_tokenizer::{DecodedText, Tokenizer as _};
use xgrammar_structural_tag::ToolChoice;
use xgrammar_structural_tag::format::Format;
use super::{BeamUnifiedParser, HfPosition, build_output_grammar, other_names};
use crate::output_grammar::test_utils::outline;
use crate::output_grammar::{GrammarCoverage, OutputGrammarContext, ToolStrictLevel};
use crate::tool::Tool;
use crate::unified::{ResponseTemplate, UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};
/// The `tag_pattern` of the checkpoint of [`template`].
const TAG_PATTERN: &str =
r#"<rfl:parameter name="(?P<key>[^"]+)">\n(?P<value>.*?)\n</rfl:parameter>"#;
/// [`TAG_PATTERN`], but a value may also end where the next parameter opens.
const LENIENT_TAG_PATTERN: &str = r#"<rfl:parameter name="(?P<key>[^"]+)">\n(?P<value>.*?)\n(?:</rfl:parameter>|(?=<rfl:parameter name="))"#;
/// The Beam-1.0 `response_template`, with `tag_pattern`.
fn template(tag_pattern: &str) -> Arc<ResponseTemplate> {
template_with(tag_pattern, json!("<|end|><|start|>assistant"))
}
/// [`template`] with `thinking_close` as the reasoning closer.
fn template_with(tag_pattern: &str, thinking_close: Value) -> Arc<ResponseTemplate> {
let template = json!({
"defaults": {"role": "assistant"},
"start_anchor_pattern": r"<\|start\|>assistant",
"fields": {
"thinking": {
"open": "<|channel|>analysis<|message|>",
"close": thinking_close,
"content": "text"
},
"content": {
"open": "<|channel|>final<|message|>",
"close": ["<|return|>", "<|end|>", "<|call|>"],
"content": "text"
},
"tool_calls": {
"open_pattern": r" to=(?P<name>[^\s<]+)<\|channel\|>analysis <\|constrain\|>xml<\|message\|>",
"close": ["<|end|><|start|>assistant", "<|call|>"],
"repeats": true,
"content": "xml-inline",
"content_args": {
"tag_pattern": tag_pattern,
"value_parser": {"name": "text", "args": {"strip": false}}
},
"transform": {
"type": "function",
"function": {"name": "{name}", "arguments": "{content}"}
}
}
}
});
Arc::new(ResponseTemplate::from_json(&template).unwrap())
}
fn tools() -> Vec<Tool> {
vec![
Tool {
name: "get_weather".to_string(),
description: None,
parameters: json!({
"type": "object",
"properties": {
"location": {"type": "string"},
"days": {"type": "integer"},
"unit": {"enum": ["c", "f"]}
},
"required": ["location"],
"additionalProperties": false
}),
strict: Some(true),
defer_loading: None,
},
Tool {
name: "get_time".to_string(),
description: None,
parameters: json!({
"type": "object",
"properties": {"zone": {"type": "string"}}
}),
strict: None,
defer_loading: None,
},
]
}
/// A parser initialized from `prompt`, the rendered conversation.
fn parser(prompt: &str) -> BeamUnifiedParser {
parser_with(TAG_PATTERN, prompt)
}
fn parser_with(tag_pattern: &str, prompt: &str) -> BeamUnifiedParser {
let tokenizer = Arc::new(TestTokenizer::new());
let prompt_ids = tokenizer.encode(prompt, false).unwrap();
let mut parser = BeamUnifiedParser::new(template(tag_pattern), &tools(), tokenizer);
parser.initialize(&prompt_ids).unwrap();
parser
}
const PROMPT: &str = "<|start|>user<|message|>Weather in Paris?<|end|><|start|>assistant";
fn grammar(position: HfPosition, tool_choice: ToolChoice, parallel: bool) -> Option<String> {
grammar_with(position, &tools(), tool_choice, parallel, false, false)
}
fn grammar_with(
position: HfPosition,
tools: &[Tool],
tool_choice: ToolChoice,
parallel: bool,
always_constrain: bool,
lenient: bool,
) -> Option<String> {
let ctx = OutputGrammarContext {
tools,
tool_choice: &tool_choice,
tool_strict_level: ToolStrictLevel::Auto,
parallel_tool_calls: parallel,
answer: None,
};
grammar_in(position, &ctx, always_constrain, lenient)
}
fn grammar_in(
position: HfPosition,
ctx: &OutputGrammarContext<'_>,
always_constrain: bool,
lenient: bool,
) -> Option<String> {
let grammar = build_output_grammar(
position,
ctx,
always_constrain,
lenient,
template(TAG_PATTERN).reasoning_closers(),
)
.unwrap()?;
assert_eq!(grammar.coverage, GrammarCoverage::FromTokenZero);
Some(outline(&grammar.format))
}
fn constrained_grammar(
position: HfPosition,
tool_choice: ToolChoice,
answer: Format,
always_constrain: bool,
) -> Option<String> {
let tools = tools();
let ctx = OutputGrammarContext {
tools: &tools,
tool_choice: &tool_choice,
tool_strict_level: ToolStrictLevel::Auto,
parallel_tool_calls: true,
answer: Some(&answer),
};
grammar_in(position, &ctx, always_constrain, false)
}
#[test]
fn position_follows_the_prompt_and_then_the_output() {
let position = |suffix: &str| parser(&format!("{PROMPT}{suffix}")).inner.position();
assert_eq!(position(""), HfPosition::ResponseStart);
assert_eq!(
position("<|channel|>analysis<|message|>Let me"),
HfPosition::Reasoning
);
// The reasoning closer ends with the anchor, so the next message starts.
assert_eq!(
position("<|channel|>analysis<|message|>Done.<|end|><|start|>assistant"),
HfPosition::ResponseStart
);
// A partial closer stays buffered.
assert_eq!(
position("<|channel|>analysis<|message|>Done.<|end|>"),
HfPosition::Other
);
assert_eq!(
position("<|channel|>final<|message|>It is"),
HfPosition::Content
);
// A partial boundary stays buffered.
assert_eq!(position("<|channel|>ana"), HfPosition::Other);
// Without the anchor, the parser parses the output from the start.
assert_eq!(
parser("no anchor").inner.position(),
HfPosition::ResponseStart
);
let mut parser = parser(PROMPT);
parser
.parse_into(
DecodedText::unattributed("<|channel|>analysis<|message|>Hmm"),
&mut UnifiedParserOutput::default(),
)
.unwrap();
assert_eq!(parser.inner.position(), HfPosition::Reasoning);
}
#[test]
fn required_calls_any_tool_after_reasoning() {
expect_test::expect![[r#"
sequence
tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant`
sequence
/[^<]/
text
sequence
or
sequence
` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>`
or
sequence
sequence
optional `\n`
tag `<rfl:parameter name="location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
sequence
optional
sequence
optional `\n`
tag `<rfl:parameter name="days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter>`
sequence
optional `\n`
tag `<rfl:parameter name="location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
sequence
` to=get_time<|channel|>analysis <|constrain|>xml<|message|>`
sequence
optional
sequence
optional `\n`
tag `<rfl:parameter name="zone">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
star
sequence
`<|end|><|start|>assistant`
or
sequence
` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>`
or
sequence
sequence
optional `\n`
tag `<rfl:parameter name="location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
sequence
optional
sequence
optional `\n`
tag `<rfl:parameter name="days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter>`
sequence
optional `\n`
tag `<rfl:parameter name="location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
sequence
` to=get_time<|channel|>analysis <|constrain|>xml<|message|>`
sequence
optional
sequence
optional `\n`
tag `<rfl:parameter name="zone">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
optional `<|call|>`
"#]].assert_eq(
&grammar(HfPosition::ResponseStart, ToolChoice::required(), true).unwrap(),
);
}
#[test]
fn auto_with_a_strict_tool_also_allows_an_answer() {
expect_test::expect![[r#"
sequence
tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant`
sequence
/[^<]/
text
or
sequence
or
sequence
` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>`
or
sequence
sequence
optional `\n`
tag `<rfl:parameter name="location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
sequence
optional
sequence
optional `\n`
tag `<rfl:parameter name="days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter>`
sequence
optional `\n`
tag `<rfl:parameter name="location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
sequence
` to=get_time<|channel|>analysis <|constrain|>xml<|message|>`
sequence
optional
sequence
optional `\n`
tag `<rfl:parameter name="zone">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
sequence
sequence
` to=`
/(?:[^\s<g][^\s<]*|g(?:[^\s<e][^\s<]*|e(?:[^\s<t][^\s<]*|t(?:[^\s<_][^\s<]*|_(?:[^\s<tw][^\s<]*|t(?:[^\s<i][^\s<]*|i(?:[^\s<m][^\s<]*|m(?:[^\s<e][^\s<]*|e(?:[^\s<][^\s<]*))?)?)?|w(?:[^\s<e][^\s<]*|e(?:[^\s<a][^\s<]*|a(?:[^\s<t][^\s<]*|t(?:[^\s<h][^\s<]*|h(?:[^\s<e][^\s<]*|e(?:[^\s<r][^\s<]*|r(?:[^\s<][^\s<]*))?)?)?)?)?)?)?)?)?)?)/
`<|channel|>analysis <|constrain|>xml<|message|>`
sequence
star
sequence
optional `\n`
sequence
`<rfl:parameter name="`
/[a-zA-Z_][a-zA-Z0-9_]*/
`">\n`
or
text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`]
json(number) excluding [`<|end|>`, `<|call|>`]
json(boolean) excluding [`<|end|>`, `<|call|>`]
json(null) excluding [`<|end|>`, `<|call|>`]
json({ ... }) excluding [`<|end|>`, `<|call|>`]
json(any[]) excluding [`<|end|>`, `<|call|>`]
`\n</rfl:parameter>`
star
sequence
`<|end|><|start|>assistant`
or
sequence
` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>`
or
sequence
sequence
optional `\n`
tag `<rfl:parameter name="location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
sequence
optional
sequence
optional `\n`
tag `<rfl:parameter name="days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter>`
sequence
optional `\n`
tag `<rfl:parameter name="location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
sequence
` to=get_time<|channel|>analysis <|constrain|>xml<|message|>`
sequence
optional
sequence
optional `\n`
tag `<rfl:parameter name="zone">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
sequence
sequence
` to=`
/(?:[^\s<g][^\s<]*|g(?:[^\s<e][^\s<]*|e(?:[^\s<t][^\s<]*|t(?:[^\s<_][^\s<]*|_(?:[^\s<tw][^\s<]*|t(?:[^\s<i][^\s<]*|i(?:[^\s<m][^\s<]*|m(?:[^\s<e][^\s<]*|e(?:[^\s<][^\s<]*))?)?)?|w(?:[^\s<e][^\s<]*|e(?:[^\s<a][^\s<]*|a(?:[^\s<t][^\s<]*|t(?:[^\s<h][^\s<]*|h(?:[^\s<e][^\s<]*|e(?:[^\s<r][^\s<]*|r(?:[^\s<][^\s<]*))?)?)?)?)?)?)?)?)?)?)/
`<|channel|>analysis <|constrain|>xml<|message|>`
sequence
star
sequence
optional `\n`
sequence
`<rfl:parameter name="`
/[a-zA-Z_][a-zA-Z0-9_]*/
`">\n`
or
text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`]
json(number) excluding [`<|end|>`, `<|call|>`]
json(boolean) excluding [`<|end|>`, `<|call|>`]
json(null) excluding [`<|end|>`, `<|call|>`]
json({ ... }) excluding [`<|end|>`, `<|call|>`]
json(any[]) excluding [`<|end|>`, `<|call|>`]
`\n</rfl:parameter>`
optional `<|call|>`
sequence
`<|channel|>final<|message|>`
text excluding [`<|end|>`, `<|return|>`]
optional `<|return|>`
"#]]
.assert_eq(&grammar(HfPosition::ResponseStart, ToolChoice::auto(), true).unwrap());
}
#[test]
fn named_choice_calls_that_tool_once_after_prompted_reasoning() {
expect_test::expect![[r#"
sequence
tag `` text `<|end|><|start|>assistant`
sequence
sequence
` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>`
or
sequence
sequence
optional `\n`
tag `<rfl:parameter name="location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
sequence
optional
sequence
optional `\n`
tag `<rfl:parameter name="days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter>`
sequence
optional `\n`
tag `<rfl:parameter name="location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
optional
sequence
optional `\n`
tag `<rfl:parameter name="unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
optional `<|call|>`
"#]].assert_eq(
&grammar(HfPosition::Reasoning, ToolChoice::function("get_weather"), true).unwrap(),
);
}
#[test]
fn no_grammar_without_a_constraining_request_or_known_position() {
let non_strict = vec![tools().pop().unwrap()];
for tools in [&[][..], &non_strict] {
assert!(
grammar_with(
HfPosition::ResponseStart,
tools,
ToolChoice::auto(),
true,
false,
false
)
.is_none()
);
}
assert!(
grammar_with(
HfPosition::Other,
&tools(),
ToolChoice::required(),
true,
true,
false
)
.is_none()
);
}
#[test]
fn none_with_tools_only_answers() {
expect_test::expect![[r#"
sequence
tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant`
sequence
/[^<]/
text
sequence
`<|channel|>final<|message|>`
text excluding [`<|end|>`, `<|return|>`]
optional `<|return|>`
"#]]
.assert_eq(&grammar(HfPosition::ResponseStart, ToolChoice::none(), true).unwrap());
}
#[test]
fn always_constrain_answers_without_tools() {
expect_test::expect![[r#"
sequence
tag `` text `<|end|><|start|>assistant`
sequence
`<|channel|>final<|message|>`
text excluding [`<|end|>`, `<|return|>`]
optional `<|return|>`
"#]]
.assert_eq(
&grammar_with(
HfPosition::Reasoning,
&[],
ToolChoice::auto(),
true,
true,
false,
)
.unwrap(),
);
}
#[test]
fn always_constrain_allows_non_strict_calls_or_an_answer() {
let non_strict = vec![tools().pop().unwrap()];
expect_test::expect![[r#"
sequence
tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant`
sequence
/[^<]/
text
or
sequence
or
sequence
` to=get_time<|channel|>analysis <|constrain|>xml<|message|>`
sequence
optional
sequence
optional `\n`
tag `<rfl:parameter name="zone">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`] `\n</rfl:parameter>`
sequence
sequence
` to=`
/(?:[^\s<g][^\s<]*|g(?:[^\s<e][^\s<]*|e(?:[^\s<t][^\s<]*|t(?:[^\s<_][^\s<]*|_(?:[^\s<t][^\s<]*|t(?:[^\s<i][^\s<]*|i(?:[^\s<m][^\s<]*|m(?:[^\s<e][^\s<]*|e(?:[^\s<][^\s<]*))?)?)?)?)?)?)?)/
`<|channel|>analysis <|constrain|>xml<|message|>`
sequence
star
sequence
optional `\n`
sequence
`<rfl:parameter name="`
/[a-zA-Z_][a-zA-Z0-9_]*/
`">\n`
or
text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`]
json(number) excluding [`<|end|>`, `<|call|>`]
json(boolean) excluding [`<|end|>`, `<|call|>`]
json(null) excluding [`<|end|>`, `<|call|>`]
json({ ... }) excluding [`<|end|>`, `<|call|>`]
json(any[]) excluding [`<|end|>`, `<|call|>`]
`\n</rfl:parameter>`
optional `<|call|>`
sequence
`<|channel|>final<|message|>`
text excluding [`<|end|>`, `<|return|>`]
optional `<|return|>`
"#]]
.assert_eq(
&grammar_with(
HfPosition::ResponseStart,
&non_strict,
ToolChoice::auto(),
false,
true,
false,
)
.unwrap(),
);
}
#[test]
fn always_constrain_finishes_an_open_answer() {
expect_test::expect![[r#"
sequence
text excluding [`<|end|>`, `<|return|>`]
optional `<|return|>`
"#]]
.assert_eq(
&grammar_with(
HfPosition::Content,
&tools(),
ToolChoice::auto(),
true,
true,
false,
)
.unwrap(),
);
assert!(
grammar_with(
HfPosition::Content,
&tools(),
ToolChoice::auto(),
true,
false,
false
)
.is_none()
);
assert!(
grammar_with(
HfPosition::Content,
&tools(),
ToolChoice::required(),
true,
true,
false
)
.is_none()
);
}
#[test]
fn answer_holds_the_requested_format() {
let json = || Format::json_schema(json!({"type": "object"}));
let answer = constrained_grammar(HfPosition::ResponseStart, ToolChoice::none(), json(), false);
expect_test::expect![[r#"
sequence
tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant`
sequence
/[^<]/
text
sequence
`<|channel|>final<|message|>`
json(object)
optional `<|return|>`
"#]]
.assert_eq(&answer.unwrap());
// Under `auto`, the answer that may replace calls holds it too.
let auto =
constrained_grammar(HfPosition::ResponseStart, ToolChoice::auto(), json(), false).unwrap();
assert!(auto.ends_with(
"`<|channel|>final<|message|>`\n json(object)\n optional `<|return|>`\n"
));
// A forced call leaves no answer to constrain.
assert_eq!(
constrained_grammar(
HfPosition::ResponseStart,
ToolChoice::required(),
json(),
false
),
grammar(HfPosition::ResponseStart, ToolChoice::required(), true)
);
}
#[test]
fn request_constraint_applies_alone_in_an_open_answer() {
let json = Format::json_schema(json!({"type": "object"}));
assert!(constrained_grammar(HfPosition::Content, ToolChoice::auto(), json, true).is_none());
}
/// Parse `output` after the prompt under `tag_pattern`, returning the calls as
/// `(name, arguments)`.
fn calls(tag_pattern: &str, output: &str) -> Vec<(String, Value)> {
let mut parser = parser_with(tag_pattern, PROMPT);
let mut events = UnifiedParserOutput::default();
parser.parse_into(DecodedText::unattributed(output), &mut events).unwrap();
events.append(parser.finish().unwrap());
let mut calls: Vec<(String, String)> = Vec::new();
for event in events.events {
if let UnifiedParserEvent::ToolCall(delta) = event {
if let Some(name) = delta.name {
calls.push((name, String::new()));
}
calls.last_mut().unwrap().1.push_str(&delta.arguments);
}
}
calls
.into_iter()
.map(|(name, arguments)| (name, serde_json::from_str(&arguments).unwrap()))
.collect()
}
/// Output in the grammar's canonical spelling parses back to the values the
/// grammar admitted.
#[test]
fn parser_reads_back_the_grammar_spelling() {
let parameter = |key: &str, value: &str| {
format!("<rfl:parameter name=\"{key}\">\n{value}\n</rfl:parameter>")
};
let call = |name: &str, parameters: &[String]| {
format!(
" to={name}<|channel|>analysis <|constrain|>xml<|message|>{}",
parameters.join("\n")
)
};
let output = format!(
"<|channel|>analysis<|message|>Two tools.<|end|><|start|>assistant{}<|end|><|start|>assistant{}",
call(
"get_weather",
&[
parameter("location", "Paris"),
parameter("days", "2"),
parameter("unit", "c")
]
),
call("get_time", &[parameter("zone", "multi\nline")]),
);
for tag_pattern in [TAG_PATTERN, LENIENT_TAG_PATTERN] {
assert_eq!(
calls(tag_pattern, &output),
vec![
(
"get_weather".to_string(),
json!({"location": "Paris", "days": 2, "unit": "c"})
),
("get_time".to_string(), json!({"zone": "multi\nline"})),
]
);
}
}
#[test]
fn values_end_before_the_next_parameter_where_the_template_looks_ahead() {
let output = " to=get_weather<|channel|>analysis <|constrain|>xml<|message|>\
<rfl:parameter name=\"days\">\n2\n\
<rfl:parameter name=\"location\">\nParis\n</rfl:parameter>";
assert!(parser_with(LENIENT_TAG_PATTERN, PROMPT).lenient);
assert_eq!(
calls(LENIENT_TAG_PATTERN, output),
vec![(
"get_weather".to_string(),
json!({"days": 2, "location": "Paris"})
)]
);
assert!(!parser(PROMPT).lenient);
}
#[test]
fn reasoning_ends_at_every_closer_of_the_template() {
let closers = [
"<|end|><|start|>assistant",
"<|constrain|><|start|>assistant",
"<|start|>assistant",
];
let tools = tools();
let tokenizer = Arc::new(TestTokenizer::new());
let mut parser = BeamUnifiedParser::new(
template_with(TAG_PATTERN, json!(closers)),
&tools,
tokenizer.clone(),
);
parser.initialize(&tokenizer.encode(PROMPT, false).unwrap()).unwrap();
let grammar = parser
.build_output_grammar(&OutputGrammarContext {
tools: &tools,
tool_choice: &ToolChoice::required(),
tool_strict_level: ToolStrictLevel::Auto,
parallel_tool_calls: false,
answer: None,
})
.unwrap()
.unwrap();
assert_eq!(
outline(&grammar.format).lines().nth(1).unwrap(),
" tag `<|channel|>analysis<|message|>` .. `<|constrain|><|start|>assistant` | \
`<|end|><|start|>assistant` | `<|start|>assistant`"
);
}
#[test]
fn lenient_arguments_come_in_declared_or_sorted_order() {
expect_test::expect![[r#"
sequence
tag `` text `<|end|><|start|>assistant`
sequence
sequence
` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>`
sequence
optional `\n`
`<rfl:parameter name="`
or
sequence
or
tag `location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`, `\n<rfl:parameter name="`] `\n</rfl:parameter>`
sequence
tag `location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`, `\n<rfl:parameter name="`] `\n</rfl:parameter><rfl:parameter name="` | `\n</rfl:parameter>\n<rfl:parameter name="` | `\n<rfl:parameter name="`
or
tag `days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter>`
sequence
optional tag `days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter><rfl:parameter name="` | `\n</rfl:parameter>\n<rfl:parameter name="` | `\n<rfl:parameter name="`
tag `unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
sequence
optional tag `days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n</rfl:parameter><rfl:parameter name="` | `\n</rfl:parameter>\n<rfl:parameter name="` | `\n<rfl:parameter name="`
or
tag `location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`, `\n<rfl:parameter name="`] `\n</rfl:parameter>`
sequence
tag `location">\n` text excluding [`<|end|>`, `<|call|>`, `</rfl:parameter>`, `\n<rfl:parameter name="`] `\n</rfl:parameter><rfl:parameter name="` | `\n</rfl:parameter>\n<rfl:parameter name="` | `\n<rfl:parameter name="`
tag `unit">\n` .. `\n</rfl:parameter>`
or
`c`
`f`
optional `<|call|>`
"#]].assert_eq(
&grammar_with(
HfPosition::Reasoning,
&tools(),
ToolChoice::function("get_weather"),
true,
false,
true,
)
.unwrap(),
);
}
#[test]
fn other_names_match_every_name_but_the_given_ones() {
let regex = regex_automata::meta::Regex::new(&format!(
"^(?:{})$",
other_names(&["get", "get_time", "a.b"])
))
.unwrap();
for name in [
"g",
"ge",
"gets",
"get_",
"get_t",
"get_times",
"x",
"a",
"a.",
"axb",
"a.bc",
] {
assert!(regex.is_match(name), "{name}");
}
for name in ["get", "get_time", "a.b", "", "get time", "get<"] {
assert!(!regex.is_match(name), "{name}");
}
}
rust/src/parser/src/unified/hf/content.rs
+77 -2
View File
@@ -6,8 +6,12 @@
//! Port of `content_parsers.py`. Each parser takes a chunk of captured text and
//! parses it into a single key in the output message dictionary.
use std::borrow::Cow;
use regex_automata::meta::Regex;
use regex_automata::util::iter::Searcher;
use regex_automata::util::syntax;
use regex_automata::{Input, Match};
use serde::Deserialize;
use serde_json::{Map, Value};
use thiserror_ext::AsReport as _;
@@ -22,6 +26,10 @@ use crate::unified::{self, parsing_failed};
const LAX_OPEN: char = '\u{1}';
const LAX_CLOSE: char = '\u{2}';
/// Name of the group that stands for the lookahead ending an `xml-inline` tag
/// pattern.
const AHEAD_GROUP: &str = "lookahead";
/// A configured content parser (`content` plus `content_args` of one field).
#[derive(Debug, Clone)]
pub(super) enum ContentParser {
@@ -66,6 +74,9 @@ pub(super) struct XmlInlineArgs {
key_group: usize,
/// Index of the `value` group, if the pattern declares one.
value_group: Option<usize>,
/// Index of the group standing for the lookahead that ends the pattern,
/// and the lookahead's own pattern, if any.
lookahead: Option<(usize, String)>,
value_parser: Option<Box<ContentParser>>,
/// Collect duplicate keys into a list.
merge_duplicates: bool,
@@ -143,7 +154,8 @@ impl ContentParser {
}
ContentKind::XmlInline => {
let spec: XmlInlineSpec = deserialize(&args, &context)?;
let tag_pattern = compile_regex(scope, &spec.tag_pattern)?;
let (pattern, lookahead) = final_lookahead(&spec.tag_pattern);
let tag_pattern = compile_regex(scope, &pattern)?;
let group_info = tag_pattern.group_info();
let Some(key_group) = group_info.to_index(Default::default(), "key") else {
return Err(invalid!(
@@ -152,10 +164,14 @@ impl ContentParser {
));
};
let value_group = group_info.to_index(Default::default(), "value");
let lookahead = group_info
.to_index(Default::default(), AHEAD_GROUP)
.zip(lookahead.map(str::to_string));
Self::XmlInline(XmlInlineArgs {
tag_pattern,
key_group,
value_group,
lookahead,
value_parser: ValueParserSpec::compile(scope, spec.value_parser)?,
merge_duplicates: spec.merge_duplicates,
})
@@ -193,6 +209,32 @@ impl ContentParser {
Self::KvLines(args) => args.parse(text),
}
}
/// The pattern of the lookahead that ends an `xml-inline` tag pattern, if
/// any.
pub fn lookahead(&self) -> Option<&str> {
match self {
Self::XmlInline(args) => args.lookahead.as_ref().map(|(_, ahead)| ahead.as_str()),
_ => None,
}
}
}
/// `pattern` with a lookahead that ends it matched as the group
/// [`AHEAD_GROUP`] instead, as the Rust regex engine has no lookahead, and
/// the lookahead's own pattern. The engine rejects any other lookahead.
fn final_lookahead(pattern: &str) -> (Cow<'_, str>, Option<&str>) {
if let Some((head, rest)) = pattern.split_once("(?=") {
// Only closing parentheses may follow a lookahead that ends the pattern.
let ahead = rest.trim_end_matches(')');
if ahead.len() < rest.len() && !ahead.contains(['(', ')']) {
return (
format!("{head}(?P<{AHEAD_GROUP}>{rest}").into(),
Some(ahead),
);
}
}
(pattern.into(), None)
}
impl ValueParserSpec {
@@ -258,7 +300,21 @@ impl XmlInlineArgs {
/// Parse shallow XML-ish tags into a dict.
fn parse(&self, text: &str) -> unified::Result<Value> {
let mut out = Map::new();
for captures in self.tag_pattern.captures_iter(text) {
let mut captures = self.tag_pattern.create_captures();
let mut tags = Searcher::new(Input::new(text));
while tags
.advance(|input| {
self.tag_pattern.search_captures(input, &mut captures);
// The lookahead's text is not part of the tag.
let ahead =
self.lookahead.as_ref().and_then(|(group, _)| captures.get_group(*group));
Ok(captures.get_match().map(|tag| match ahead {
Some(ahead) => Match::new(tag.pattern(), tag.start()..ahead.start),
None => tag,
}))
})
.is_some()
{
let key =
captures.get_group(self.key_group).map(|span| &text[span.range()]).ok_or_else(
|| parsing_failed!("xml-inline: named group 'key' did not participate"),
@@ -490,6 +546,25 @@ mod tests {
);
}
#[test]
fn xml_inline_tags_may_end_where_the_pattern_looks_ahead() {
let xml = parser(
"xml-inline",
json!({
"tag_pattern": r#"<arg name="(?P<key>\w+)">(?P<value>.*?)\n(?:</arg>|(?=<arg name="))"#,
}),
);
let value = xml
.parse("<arg name=\"a\">1\n<arg name=\"b\">x <arg name=\"c\"> y\n</arg>")
.unwrap();
assert_eq!(value, json!({"a": "1", "b": "x <arg name=\"c\"> y"}));
assert_eq!(xml.lookahead(), Some("<arg name=\""));
// Any other lookahead is unsupported.
let args = json!({"tag_pattern": r"(?P<key>\w+)(?=:)(?P<value>.*)"});
assert!(ContentParser::new("test", kind("xml-inline"), args.as_object().unwrap()).is_err());
}
#[test]
fn kv_lines_split_and_strip() {
let kv = parser("kv-lines", json!({"value_parser": {"name": "int"}}));
rust/src/parser/src/unified/hf/mod.rs
+1 -1
View File
@@ -25,7 +25,7 @@ mod transform;
#[cfg(test)]
mod tests;
pub use parser::HfUnifiedParser;
pub use parser::{HfPosition, HfUnifiedParser};
pub use template::ResponseTemplate;
use thiserror::Error;
use thiserror_ext::Macro;
rust/src/parser/src/unified/hf/parser.rs
+44
View File
@@ -59,6 +59,20 @@ pub struct HfUnifiedParser {
emitted_tool_count: usize,
}
/// Where the parser is in the response; see [`HfUnifiedParser::position`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HfPosition {
/// Nothing parsed since the start anchor, or since the start of the stream
/// when there is none.
ResponseStart,
/// Inside an open reasoning field.
Reasoning,
/// Inside an open content field.
Content,
/// Anywhere else.
Other,
}
/// Parser mode: the region currently receiving text.
enum HfMode {
/// No implicit field is declared: text outside regions is discarded.
@@ -156,6 +170,36 @@ impl HfUnifiedParser {
parser
}
/// Where the parser is in the response. Between
/// [`UnifiedParser::initialize`] and the first generated text, this is where
/// the prompt leaves the response.
///
/// Text outside every region, which a template without an implicit field
/// discards, leaves no state behind and so does not move the parser from
/// [`HfPosition::ResponseStart`].
pub fn position(&self) -> HfPosition {
if !self.buffer.text.is_empty() {
return HfPosition::Other;
}
match &self.mode {
HfMode::Implicit(occurrence) | HfMode::Explicit(occurrence) if occurrence.opened => {
match &self.template.regions[occurrence.region].kind {
RegionKind::Text(TextRegion {
role: TextRole::Reasoning,
..
}) => HfPosition::Reasoning,
RegionKind::Text(TextRegion {
role: TextRole::Content,
..
}) => HfPosition::Content,
_ => HfPosition::Other,
}
}
_ if self.closed.contains(&true) || self.emitted_tool_count > 0 => HfPosition::Other,
_ => HfPosition::ResponseStart,
}
}
/// Clear all per-stream state.
fn reset_state(&mut self) {
self.buffer.clear();
rust/src/parser/src/unified/hf/template.rs
+25
View File
@@ -191,6 +191,31 @@ impl ResponseTemplate {
})
}
/// The pattern of the lookahead that ends the `xml-inline` tag pattern of
/// tool-call arguments, such as for the opener of the next argument.
pub fn argument_lookahead(&self) -> Option<&str> {
self.regions.iter().find_map(|region| match &region.kind {
RegionKind::ToolCalls(calls) => calls.content.lookahead(),
_ => None,
})
}
/// The literal closers of the reasoning field, longest first, or none if it
/// has no literal close.
pub fn reasoning_closers(&self) -> &[String] {
self.regions
.iter()
.find_map(|region| match (&region.kind, &region.close) {
(RegionKind::Text(text), Some(Boundary::Literals(closers)))
if text.role == TextRole::Reasoning =>
{
Some(closers.as_slice())
}
_ => None,
})
.unwrap_or_default()
}
/// Return the start and end offsets of the last start-anchor match in
/// `prompt`.
pub(super) fn last_start_anchor(&self, prompt: &str) -> Option<(usize, usize)> {
rust/src/parser/src/unified/mod.rs
+3 -1
View File
@@ -3,6 +3,7 @@
//! Unified parser interface for reasoning and tool-call deltas.
mod beam;
mod combined;
mod gemma4;
mod hf;
@@ -10,9 +11,10 @@ mod hy;
mod inkling;
mod kimi_k3;
pub use beam::BeamUnifiedParser;
pub use combined::CombinedParser;
pub use gemma4::Gemma4UnifiedParser;
pub use hf::{HfTemplateError, HfUnifiedParser, ResponseTemplate};
pub use hf::{HfPosition, HfTemplateError, HfUnifiedParser, ResponseTemplate};
pub use hy::{HyV3UnifiedParser, HyV4UnifiedParser};
pub use inkling::InklingUnifiedParser;
pub use kimi_k3::KimiK3UnifiedParser;
tests/benchmarks/test_forward_cli.py
+13
View File
@@ -9,6 +9,7 @@ import numpy as np
import pytest
from vllm.benchmarks.forward.batch import build_batch
from vllm.benchmarks.forward.main import run_case
from vllm.benchmarks.forward.shapes import jitter_kv_lens, parse_case
from vllm.platforms import current_platform
@@ -84,6 +85,18 @@ def test_jitter_kv_lens():
jitter_kv_lens(case, 0, 1000, 0)
@pytest.mark.parametrize("capacity", [0, 2])
def test_bench_forward_rejects_oversized_verify(capacity):
"""Reject drafts beyond the runner's buffer before launching CUDA kernels."""
runner = SimpleNamespace(num_speculative_steps=capacity)
case = parse_case("verify=2*4@64:3")
with pytest.raises(
ValueError,
match=rf"Case 'verify' needs 3 draft tokens.*supports {capacity}",
):
run_case(runner, case, SimpleNamespace())
@pytest.mark.benchmark
# TODO(woosuk): Support ROCm and re-enable this test there.
@pytest.mark.skipif(
tests/compile/passes/test_fusion.py
+16
View File
@@ -107,6 +107,8 @@ AITER_KERNEL_GROUPSHAPE_COMBINATIONS = [
# Per-token with ChannelWiseTorchFP8ScaledMMLinearKernel
(ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, True),
(ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, False),
# Static per-tensor: only the aiter quant op is matched by the pass
(ROCmFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR, True),
# Blockwise
(AiterFp8BlockScaledMMKernel, GroupShape(1, 128), True),
]
@@ -197,6 +199,9 @@ class TestModel(torch.nn.Module):
# Blockwise path
if self.use_aiter_fusion and self.use_aiter_quant_op:
return [rocm_aiter_ops.get_group_quant_op()]
elif self.group_shape.is_per_tensor():
if self.use_aiter_quant_op:
return [rocm_aiter_ops.get_per_tensor_quant_op()]
else:
if self.use_aiter_quant_op:
return [rocm_aiter_ops.get_per_token_quant_op()]
@@ -221,6 +226,17 @@ class TestModel(torch.nn.Module):
AiterFusedAddRMSFp8GroupQuantPattern.FUSED_OP,
AiterRMSFp8GroupQuantPattern.FUSED_OP,
]
elif self.group_shape.is_per_tensor():
# Static per-tensor aiter fusion
from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
AiterFusedAddRMSNormStaticQuantPattern,
AiterRMSNormStaticQuantPattern,
)
return [
AiterFusedAddRMSNormStaticQuantPattern.FUSED_OP,
AiterRMSNormStaticQuantPattern.FUSED_OP,
]
else:
# Per-token aiter fusion
from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
tests/distributed/test_weight_transfer_m2n.py
+1063 -160
View File
File diff suppressed because it is too large Load Diff
tests/entrypoints/anthropic/test_anthropic_messages_conversion.py
+56 -29
View File
@@ -49,6 +49,9 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
ChatMessage,
)
from vllm.entrypoints.serve.engine.protocol import PromptTokenUsageInfo, UsageInfo
from vllm.entrypoints.serve.exception_handling.error_response import (
create_error_response,
)
from vllm.entrypoints.serve.exception_handling.handlers.validation import (
validation_exception_handler,
)
@@ -1034,7 +1037,7 @@ def _make_stream_chunk(
stop_reason: int | str | None = None,
choices: list[ChatCompletionResponseStreamChoice] | None = None,
usage: UsageInfo | None = None,
) -> str:
) -> ChatCompletionStreamResponse:
if choices is None:
choices = [
ChatCompletionResponseStreamChoice(
@@ -1051,7 +1054,7 @@ def _make_stream_chunk(
choices=choices,
usage=usage,
)
return f"data: {chunk.model_dump_json()}"
return chunk
def _tc(*, args, id=None, name=None):
@@ -1062,6 +1065,39 @@ def _tc(*, args, id=None, name=None):
)
@pytest.mark.asyncio
@pytest.mark.parametrize("block_type", [None, "text", "thinking", "tool_use"])
async def test_stream_error_preserves_message_and_closes_block(block_type):
async def source():
if block_type is not None:
yield _make_stream_chunk(delta=DeltaMessage(role="assistant"))
if block_type == "text":
delta = DeltaMessage(content="Hello")
elif block_type == "thinking":
delta = DeltaMessage(reasoning="Thinking")
else:
delta = DeltaMessage(
tool_calls=[_tc(id="call_1", name="search", args='{"q":')]
)
yield _make_stream_chunk(delta=delta)
yield create_error_response(ValueError("backend failed"))
converter = _make_stream_converter()
events = _parse_sse_events(
[event async for event in converter.message_stream_converter(source())]
)
assert events[-1] == (
"error",
{
"type": "error",
"error": {"type": "internal_error", "message": "backend failed"},
},
)
assert not any(event_type == "message_stop" for event_type, _ in events)
if block_type is not None:
assert events[-2][0] == "content_block_stop"
class TestMessageStreamConverterToolUseContentBuffering:
"""Regression test for tool_use arguments being silently dropped.
@@ -1074,7 +1110,7 @@ class TestMessageStreamConverterToolUseContentBuffering:
async def test_tool_use_args_not_dropped_when_content_in_same_chunk(
self,
):
async def sse_input():
async def chunk_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant"),
usage=UsageInfo(prompt_tokens=10, total_tokens=10),
@@ -1111,11 +1147,10 @@ class TestMessageStreamConverterToolUseContentBuffering:
completion_tokens=20,
),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
async for event in converter.message_stream_converter(chunk_input()):
output.append(event)
events = _parse_sse_events(output)
@@ -1156,7 +1191,7 @@ class TestMessageStreamConverterToolUseContentBuffering:
stream jumps straight from finish_reason to [DONE], skipping the
empty-choices usage chunk."""
async def sse_input():
async def chunk_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant"),
usage=UsageInfo(prompt_tokens=10, total_tokens=10),
@@ -1178,11 +1213,10 @@ class TestMessageStreamConverterToolUseContentBuffering:
finish_reason="tool_calls",
)
# No empty-choices usage chunk — go straight to [DONE].
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
async for event in converter.message_stream_converter(chunk_input()):
output.append(event)
events = _parse_sse_events(output)
@@ -1216,7 +1250,7 @@ class TestMessageStartIncludesTypeAndRole:
@pytest.mark.asyncio
async def test_message_start_contains_message_type_and_role(self):
async def sse_input():
async def chunk_input():
yield _make_stream_chunk(
delta=DeltaMessage(content="Hello"),
usage=UsageInfo(
@@ -1226,11 +1260,10 @@ class TestMessageStartIncludesTypeAndRole:
),
)
yield _make_stream_chunk(finish_reason="stop")
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
async for event in converter.message_stream_converter(chunk_input()):
output.append(event)
events = _parse_sse_events(output)
@@ -1260,7 +1293,7 @@ class TestStreamingCacheUsageSemantics:
message_start omits cache fields. The final chunk carries
prompt_tokens_details, so message_delta carries resolved values."""
async def sse_input():
async def chunk_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant", content="hi"),
usage=UsageInfo(prompt_tokens=100, total_tokens=100),
@@ -1277,11 +1310,10 @@ class TestStreamingCacheUsageSemantics:
),
),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
async for event in converter.message_stream_converter(chunk_input()):
output.append(event)
events = _parse_sse_events(output)
@@ -1305,7 +1337,7 @@ class TestStreamingCacheUsageSemantics:
"""When the final chunk reports cached_tokens=0, message_delta carries
cache fields = 0 (cache miss); message_start still omits them."""
async def sse_input():
async def chunk_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant"),
usage=UsageInfo(prompt_tokens=50, total_tokens=50),
@@ -1322,11 +1354,10 @@ class TestStreamingCacheUsageSemantics:
),
),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
async for event in converter.message_stream_converter(chunk_input()):
output.append(event)
events = _parse_sse_events(output)
@@ -1346,7 +1377,7 @@ class TestStreamingCacheUsageSemantics:
"""If --enable-prompt-tokens-details is off, no chunk carries cache
info; both message_start and message_delta omit cache fields."""
async def sse_input():
async def chunk_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant"),
usage=UsageInfo(prompt_tokens=30, total_tokens=30),
@@ -1356,11 +1387,10 @@ class TestStreamingCacheUsageSemantics:
choices=[],
usage=UsageInfo(prompt_tokens=30, completion_tokens=2, total_tokens=32),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
async for event in converter.message_stream_converter(chunk_input()):
output.append(event)
events = _parse_sse_events(output)
@@ -1628,7 +1658,7 @@ class TestStopSequenceReason:
@pytest.mark.asyncio
async def test_streaming_stop_string_maps_to_stop_sequence(self):
async def sse_input():
async def chunk_input():
yield _make_stream_chunk(
delta=DeltaMessage(content="hi"),
usage=UsageInfo(prompt_tokens=5, total_tokens=5, completion_tokens=0),
@@ -1638,11 +1668,10 @@ class TestStopSequenceReason:
choices=[],
usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
async for event in converter.message_stream_converter(chunk_input()):
output.append(event)
events = _parse_sse_events(output)
@@ -1654,7 +1683,7 @@ class TestStopSequenceReason:
async def test_streaming_no_stop_string_emits_explicit_null_stop_sequence(self):
"""exclude_unset=True drops stop_sequence unless it is set explicitly."""
async def sse_input():
async def chunk_input():
yield _make_stream_chunk(delta=DeltaMessage(role="assistant"))
yield _make_stream_chunk(delta=DeltaMessage(content="hi"))
yield _make_stream_chunk(finish_reason="stop")
@@ -1662,11 +1691,10 @@ class TestStopSequenceReason:
choices=[],
usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
async for event in converter.message_stream_converter(chunk_input()):
output.append(event)
events = _parse_sse_events(output)
@@ -1724,7 +1752,7 @@ class TestToolUseStopReason:
],
)
async def test_streaming(self, args, finish_reason, stop_reason, expected):
async def sse_input():
async def chunk_input():
yield _make_stream_chunk(delta=DeltaMessage(role="assistant"))
yield _make_stream_chunk(
delta=DeltaMessage(
@@ -1738,11 +1766,10 @@ class TestToolUseStopReason:
choices=[],
usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3),
)
yield "data: [DONE]"
converter = _make_stream_converter()
events = _parse_sse_events(
[e async for e in converter.message_stream_converter(sse_input())]
[e async for e in converter.message_stream_converter(chunk_input())]
)
msg_deltas = [data for ev_type, data in events if ev_type == "message_delta"]
tests/entrypoints/cohere/test_serving_streaming.py
+60 -75
View File
@@ -5,7 +5,7 @@
The stream-translation entry point is
:meth:`CohereServingChatV2._chat_completion_stream_to_v2`, which turns an
async iterable of OpenAI SSE chunks into Cohere's
async iterable of OpenAI chunk objects into Cohere's
``message-start → (content|tool-call|citation)* → message-end → [DONE]``
event stream.
@@ -15,7 +15,7 @@ state machine.
"""
import json
from collections.abc import AsyncGenerator
from collections.abc import AsyncGenerator, Sequence
from typing import Any
import pytest
@@ -26,10 +26,7 @@ from vllm.entrypoints.cohere.cohere_chat_message import (
from vllm.entrypoints.cohere.cohere_chat_message import (
CitationSource,
)
from vllm.entrypoints.cohere.protocol import (
CohereChatV2Request,
MessageStartEvent,
)
from vllm.entrypoints.cohere.protocol import MessageStartEvent
from vllm.entrypoints.cohere.serving import (
_DONE_FRAME,
CohereServingChatV2,
@@ -41,6 +38,7 @@ from vllm.entrypoints.cohere.serving import (
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionStreamResponse,
)
from vllm.entrypoints.serve.engine.protocol import ErrorResponse
# ----------------------------------------------------------------------
# Helpers
@@ -85,9 +83,8 @@ def _make_chunk(
usage: dict[str, Any] | None = None,
omit_choices: bool = False,
citations: list[Any] | None = None,
) -> str:
"""Build the ``data: {...}\\n\\n`` SSE frame the production code
consumes."""
) -> ChatCompletionStreamResponse:
"""Build the raw chat chunk the production converter consumes."""
delta: dict[str, Any] = {}
if role is not None:
delta["role"] = role
@@ -119,27 +116,23 @@ def _make_chunk(
payload["choices"] = []
if usage is not None:
payload["usage"] = usage
# Use ChatCompletionStreamResponse to normalize the payload shape.
chunk = ChatCompletionStreamResponse.model_validate(payload)
return f"data: {chunk.model_dump_json(exclude_none=False)}\n\n"
return ChatCompletionStreamResponse.model_validate(payload)
def _make_done() -> str:
return "data: [DONE]\n\n"
async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]:
async def _stream_from(
items: Sequence[ChatCompletionStreamResponse | ErrorResponse],
) -> AsyncGenerator[ChatCompletionStreamResponse | ErrorResponse, None]:
for item in items:
yield item
async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]:
async def _drain(
serving: CohereServingChatV2,
items: Sequence[ChatCompletionStreamResponse | ErrorResponse],
) -> list[str]:
"""Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect
the emitted SSE frames."""
request = CohereChatV2Request(
model="m", messages=[{"role": "user", "content": "hi"}], stream=True
)
gen = serving._chat_completion_stream_to_v2(_stream_from(items), request)
gen = serving._chat_completion_stream_to_v2(_stream_from(items))
return [frame async for frame in gen]
@@ -793,15 +786,14 @@ class TestChatCompletionStreamToV2:
assert "content-end" not in types
@pytest.mark.asyncio
async def test_done_marker_in_middle_closes_open_block(self):
# Some upstreams send [DONE] without a trailing usage-only chunk.
async def test_stream_exhaustion_closes_open_block_without_usage(self):
# The raw stream may finish without a trailing usage-only chunk.
# The translator must still emit message-end before [DONE] so
# Cohere clients don't hang.
serving = _serving(is_reasoning_model=True)
items = [
_make_chunk(role="assistant"),
_make_chunk(content="Hi"),
_make_done(),
]
frames = await _drain(serving, items)
types = [_parse_event(f)["type"] for f in frames[:-1]]
@@ -815,72 +807,65 @@ class TestChatCompletionStreamToV2:
assert frames[-1] == _DONE_FRAME
@pytest.mark.asyncio
async def test_skips_empty_and_non_data_lines(self):
@pytest.mark.parametrize("block_type", [None, "text", "thinking", "tool_call"])
async def test_upstream_error_preserves_message(self, block_type):
serving = _serving()
items = [
"\n",
"event: ping\n\n",
_make_chunk(role="assistant"),
"data: \n\n", # empty data
_make_chunk(content="Hi"),
_make_chunk(
omit_choices=True,
usage={
"prompt_tokens": 1,
"completion_tokens": 1,
"total_tokens": 2,
},
),
]
items: list[ChatCompletionStreamResponse | ErrorResponse] = []
if block_type is not None:
items.append(_make_chunk(role="assistant"))
if block_type == "text":
items.append(_make_chunk(content="Hi"))
elif block_type == "thinking":
items.append(_make_chunk(reasoning="Thinking"))
else:
items.append(
_make_chunk(
tool_calls=[
{
"index": 0,
"id": "call_1",
"function": {"name": "search", "arguments": '{"q":'},
}
]
)
)
items.append(serving.create_error_response(ValueError("backend failed")))
frames = await _drain(serving, items)
assert frames[-1] == _DONE_FRAME
# Should still produce a complete lifecycle.
types = [_parse_event(f)["type"] for f in frames[:-1]]
assert types[0] == "message-start"
assert types[-1] == "message-end"
events = [_parse_event(frame) for frame in frames[:-1]]
error_ends = [event for event in events if event["type"] == "message-end"]
assert len(error_ends) == 1
assert error_ends[0]["delta"] == {
"finish_reason": "ERROR",
"error": "backend failed",
}
assert events[-1] == error_ends[0]
if block_type is not None:
assert events[-2]["type"] == (
"tool-call-end" if block_type == "tool_call" else "content-end"
)
@pytest.mark.asyncio
async def test_exception_in_chunk_parsing_emits_error_message_end(self):
# An invalid JSON payload after the first chunk triggers the
# error path: a synthetic message-end with finish_reason=ERROR
# followed by [DONE].
serving = _serving()
items = [
_make_chunk(role="assistant"),
"data: {not valid json}\n\n",
]
frames = await _drain(serving, items)
assert frames[-1] == _DONE_FRAME
# Find the error-shaped message-end.
error_end = _parse_event(frames[-2])
assert error_end["type"] == "message-end"
assert error_end["delta"]["finish_reason"] == "ERROR"
assert "error" in error_end["delta"]
@pytest.mark.asyncio
async def test_error_message_end_is_sanitized(self):
@pytest.mark.parametrize("started", [False, True])
async def test_error_message_end_is_sanitized(self, started):
# The client-visible ``error`` field on the terminal message-end
# must be routed through ``sanitize_message`` so tracebacks,
# host filesystem paths, and Python object memory addresses
# can't leak from a mid-stream exception into the SSE payload.
serving = _serving()
async def _raises_mid_stream() -> AsyncGenerator[str, None]:
yield _make_chunk(role="assistant")
async def _raises_mid_stream() -> AsyncGenerator[
ChatCompletionStreamResponse | ErrorResponse, None
]:
if started:
yield _make_chunk(role="assistant")
raise RuntimeError(
"boom <Foo at 0x7fabc0deface> reading /Users/dev/vllm/x.py"
)
request = CohereChatV2Request(
model="m",
messages=[{"role": "user", "content": "hi"}],
stream=True,
)
frames = [
f
async for f in serving._chat_completion_stream_to_v2(
_raises_mid_stream(), request
)
f async for f in serving._chat_completion_stream_to_v2(_raises_mid_stream())
]
assert frames[-1] == _DONE_FRAME
error_end = _parse_event(frames[-2])
tests/entrypoints/openai/chat_completion/test_serving_chat.py
+154 -1
View File
@@ -28,6 +28,7 @@ from vllm.entrypoints.generate.base.serving import build_per_request_timing_metr
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
ChatCompletionResponse,
ChatCompletionStreamResponse,
)
from vllm.entrypoints.openai.chat_completion.serving import (
OpenAIServingChat,
@@ -43,7 +44,7 @@ from vllm.entrypoints.openai.models.serving import (
from vllm.entrypoints.openai.parser.harmony_utils import get_encoding
from vllm.entrypoints.serve.engine.protocol import ErrorResponse
from vllm.entrypoints.serve.utils.tool_calls_utils import resolve_finish_reason
from vllm.exceptions import QueueOverflowError, VLLMValidationError
from vllm.exceptions import GenerationError, QueueOverflowError, VLLMValidationError
from vllm.inputs import TokensPrompt
from vllm.logprobs import Logprob
from vllm.multimodal.inputs import PlaceholderRange
@@ -716,6 +717,149 @@ async def _collect_metrics_stream_chunks(
return chunks
@pytest.mark.asyncio
@pytest.mark.parametrize("raw", [False, True])
async def test_streaming_chunks_keep_usage_and_metadata(raw):
serving = _build_minimal_metrics_serving_chat(False)
request = ChatCompletionRequest(
model="test-model",
messages=[{"role": "user", "content": "Test"}],
stream=True,
stream_options={"include_usage": True},
)
metadata = RequestResponseMetadata(request_id="chatcmpl-test")
method = (
serving.chat_completion_stream_generator_raw
if raw
else serving.chat_completion_stream_generator
)
chunks = [
chunk
async for chunk in method(
request,
_single_request_output(_make_metrics_request_output()),
"chatcmpl-test",
"test-model",
conversation=[],
tokenizer=MagicMock(),
request_metadata=metadata,
)
]
if raw:
assert all(isinstance(chunk, ChatCompletionStreamResponse) for chunk in chunks)
else:
assert chunks[-1] == "data: [DONE]\n\n"
first = json.loads(chunks[0].removeprefix("data: "))
assert first["choices"][0]["logprobs"] is None
usage = json.loads(chunks[-2].removeprefix("data: "))
assert "system_fingerprint" not in usage
assert "prompt_tokens_details" not in usage["usage"]
chunks = [
ChatCompletionStreamResponse.model_validate_json(
chunk.removeprefix("data: ")
)
for chunk in chunks[:-1]
]
assert len(chunks) == 3
assert chunks[0].choices[0].delta.role == "assistant"
assert chunks[1].choices[0].delta.content == "Hello"
assert chunks[1].choices[0].finish_reason == "stop"
assert chunks[2].choices == []
assert chunks[2].usage.prompt_tokens == 3
assert chunks[2].usage.completion_tokens == 2
assert metadata.final_usage_info.total_tokens == 5
@pytest.mark.asyncio
async def test_streaming_serialization_failure_emits_error_and_done():
serving = _build_minimal_metrics_serving_chat(False)
serving.parser_cls = HarmonyParser
parser = MagicMock()
parser.parse_delta.return_value = DeltaMessage(
content="Hello", extra_object=object()
)
serving._make_parser = MagicMock(return_value=parser)
request = ChatCompletionRequest(
model="test-model", messages=[{"role": "user", "content": "Test"}], stream=True
)
frames = [
frame
async for frame in serving.chat_completion_stream_generator(
request,
_single_request_output(_make_metrics_request_output()),
"chatcmpl-test",
"test-model",
conversation=[],
tokenizer=MagicMock(),
request_metadata=RequestResponseMetadata(request_id="chatcmpl-test"),
)
]
assert frames[-1] == "data: [DONE]\n\n"
error = ErrorResponse.model_validate_json(frames[-2].removeprefix("data: "))
assert "serialize" in error.error.message
@pytest.mark.asyncio
@pytest.mark.parametrize("raw", [False, True])
@pytest.mark.parametrize(
"failure", ["parser", "before_first", "after_content", "generation"]
)
async def test_streaming_errors_preserve_envelope(raw, failure):
serving = _build_minimal_metrics_serving_chat(False)
error = (
GenerationError("backend failed")
if failure == "generation"
else ValueError("backend failed")
)
if failure == "parser":
serving.parser_cls = HarmonyParser
serving._make_parser = MagicMock(side_effect=error)
async def results():
if failure == "after_content":
output = _make_metrics_request_output()
output.outputs[0].finish_reason = None
output.finished = False
yield output
raise error
request = ChatCompletionRequest(
model="test-model", messages=[{"role": "user", "content": "Test"}], stream=True
)
method = (
serving.chat_completion_stream_generator_raw
if raw
else serving.chat_completion_stream_generator
)
chunks = [
chunk
async for chunk in method(
request,
results(),
"chatcmpl-test",
"test-model",
conversation=[],
tokenizer=MagicMock(),
request_metadata=RequestResponseMetadata(request_id="chatcmpl-test"),
)
]
if raw:
assert all(
isinstance(chunk, (ChatCompletionStreamResponse, ErrorResponse))
for chunk in chunks
)
assert isinstance(chunks[-1], ErrorResponse)
error_response = chunks[-1]
else:
assert chunks[-1] == "data: [DONE]\n\n"
assert chunks.count("data: [DONE]\n\n") == 1
error_response = ErrorResponse.model_validate_json(
chunks[-2].removeprefix("data: ")
)
assert error_response.error.message == "backend failed"
assert error_response.error.code == (500 if failure == "generation" else 400)
def test_build_per_request_timing_metrics_valid_timestamps():
metrics = build_per_request_timing_metrics(
_PER_REQUEST_STATS, num_generation_tokens=10
@@ -1188,6 +1332,15 @@ async def test_streaming_reasoning_usage_lifecycle_edges(
await iterator.aclose()
error_chunks = [c for c in chunks if isinstance(c, dict) and "error" in c]
if name == "empty-finish":
# A choice that finishes without output still gets its last chunk, unparsed.
assert any(
choice.get("finish_reason") == "stop"
for item in chunks
if isinstance(item, dict)
for choice in item.get("choices", [])
)
parsers[0].parse_delta.assert_not_called()
call_count = sum(parser.count_reasoning_tokens.call_count for parser in parsers)
if expected_tokens is not None:
assert outcome == "exhausted" and chunks[-1] == "[DONE]" and not error_chunks
tests/entrypoints/openai/completion/test_completion_error.py
+72 -1
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from dataclasses import dataclass, field
from typing import Any
from unittest.mock import AsyncMock, MagicMock
@@ -17,7 +18,7 @@ from vllm.entrypoints.scale_out.render.serving import ServingRender
from vllm.entrypoints.serve.engine.protocol import ErrorResponse
from vllm.exceptions import GenerationError, VLLMValidationError
from vllm.logprobs import Logprob
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.outputs import CompletionOutput, RequestOutput, WeightVersionSpan
from vllm.renderers.deepseek_v4 import DeepseekV4Renderer
from vllm.renderers.hf import HfRenderer
from vllm.renderers.online_renderer import OnlineRenderer
@@ -310,6 +311,76 @@ def test_completion_metrics_carries_both_timing_and_spec_decode():
assert response.metrics.speculative_decoding.num_spec_steps == 2
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
@pytest.mark.parametrize("opt_in", [False, True])
async def test_completion_weight_versions(stream: bool, opt_in: bool):
"""Each choice gets its spans on its final output, even with no tokens."""
expected = [{"version": "v1", "start": 0, "end": 2}]
spans = [WeightVersionSpan(**span) for span in expected]
def output(index, token_ids, finish_reason=None, weight_versions=None):
return CompletionOutput(
index=index,
text="x" * len(token_ids),
token_ids=token_ids,
cumulative_logprob=None,
logprobs=None,
finish_reason=finish_reason,
weight_versions=weight_versions,
)
# Choice 1 is aborted before its first token.
final = [output(1, [], "abort", [])]
if stream:
steps = [[output(0, [100])], [output(0, [101], "length", spans), *final]]
else:
steps = [[output(0, [100, 101], "length", spans), *final]]
async def mock_generate(*args, **kwargs):
for i, outputs in enumerate(steps):
yield RequestOutput(
request_id="test-id",
prompt="Test prompt",
prompt_token_ids=[1, 2, 3],
prompt_logprobs=None,
outputs=outputs,
finished=i == len(steps) - 1,
)
mock_engine = MagicMock(spec=AsyncLLM)
mock_engine.errored = False
mock_engine.model_config = MockModelConfig()
mock_engine.input_processor = MagicMock()
mock_engine.renderer = _build_renderer(mock_engine.model_config)
mock_engine.generate = MagicMock(side_effect=mock_generate)
request = CompletionRequest(
model=MODEL_NAME,
prompt="Test prompt",
max_tokens=10,
n=2,
stream=stream,
return_weight_versions=opt_in,
)
response = await _build_serving_completion(mock_engine).create_completion(request)
if stream:
finals = {}
async for line in response:
payload = line.removeprefix("data: ").strip()
if payload == "[DONE]":
continue
for choice in json.loads(payload)["choices"]:
if choice.get("finish_reason"):
finals[choice["index"]] = choice
else:
assert choice.get("weight_versions") is None
got = [finals[i].get("weight_versions") for i in (0, 1)]
else:
got = [choice["weight_versions"] for choice in response.model_dump()["choices"]]
assert got == ([expected, []] if opt_in else [None, None])
def _build_online_renderer_for_completion(
model_config: MockModelConfig,
renderer: Any | None = None,
tests/entrypoints/pooling/scoring/test_io_processor_unit.py
+46 -1
View File
@@ -10,7 +10,7 @@ import torch
from tokenizers import Tokenizer, models, pre_tokenizers, processors
from transformers import TokenizersBackend
from vllm import TokensPrompt
from vllm import PoolingParams, TokensPrompt
from vllm.entrypoints.pooling.scoring.io_processor import (
BiEncoderIOProcessor,
CrossEncoderIOProcessor,
@@ -18,7 +18,9 @@ from vllm.entrypoints.pooling.scoring.io_processor import (
LateInteractionIOProcessor,
_apply_post_tokenization_to_token_type_ids,
)
from vllm.entrypoints.pooling.scoring.typing import ScoringData
from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids
from vllm.entrypoints.pooling.typing import OfflineScoringInputsContext
from vllm.outputs import PoolingOutput, PoolingRequestOutput, RequestError
from vllm.renderers import TokenizeParams
@@ -81,6 +83,49 @@ def test_llm_reranker_tokenization_is_independent_of_nonbinding_doc_limit(
assert nonbinding["prompt_token_ids"] == expected
@pytest.mark.parametrize(
"default_template,request_template,expected",
[
("{{ messages[0].content }} {{ messages[1].content }}", None, "a bc"),
(
"{{ messages[0].content }} {{ messages[1].content }}",
"{{ messages[1].content }} {{ messages[0].content }}",
"bc a",
),
(None, "{{ messages[1].content }} {{ messages[0].content }}", "bc a"),
(None, None, "abc"),
],
)
def test_offline_scoring_uses_default_template_unless_overridden(
llm_reranker_processor, default_template, request_template, expected
):
processor = llm_reranker_processor
processor.chat_template = default_template
processor.renderer = SimpleNamespace(
default_cmpl_tok_params=TokenizeParams(max_total_tokens=32)
)
ctx = OfflineScoringInputsContext(
pooling_task="classify",
tokenization_kwargs=None,
lora_request=None,
priorities=None,
scoring_data=ScoringData(data_1=["a"], data_2=["bc"]),
chat_template=request_template,
pooling_params=PoolingParams(),
)
factory, _ = processor.get_request_factory_offline(ctx)
(request,) = factory()
text, prompt = processor.get_score_prompt(
request["data_1"],
request["data_2"],
request["tok_params"].get_encode_kwargs(),
chat_template=request["chat_template"],
)
assert text == expected
assert prompt["prompt_token_ids"] == processor.tokenizer.encode(expected)
def test_token_type_ids_stay_aligned_with_a_truncated_padded_prompt():
"""The cross-encoder segment boundary must survive truncate + pad.
tests/entrypoints/scale_out/derender/test_derender.py
+100 -12
View File
@@ -10,6 +10,7 @@ import pytest
import pytest_asyncio
from tests.utils import RemoteLaunchRenderServer
from vllm.entrypoints.generate.base.protocol import FunctionCall
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
GenerateTokensResponse,
@@ -1133,6 +1134,50 @@ async def test_e2e_parsed_reasoning(parser_client, parser_tokenizer):
assert "<think>" not in msg["content"]
@pytest.mark.asyncio
@pytest.mark.parametrize("include_reasoning", [True, False])
async def test_e2e_parsed_logprobs_follow_include_reasoning(
parser_client, parser_tokenizer, include_reasoning
):
"""With a parser configured, batch derender returns logprobs only when
reasoning is shown. Hidden reasoning drops them, as the coupled chat
endpoint does, because decoded logprob tokens would reveal it."""
messages = [{"role": "user", "content": "What is 2+3?"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
output_ids = _require_markers_survive(
parser_tokenizer, "<think>Two plus three.</think>Five.", "</think>"
)
generate_response = _e2e_generate_response(output_ids)
generate_response["choices"][0]["logprobs"] = {
"content": [
{"token_id": tid, "logprob": -0.5, "rank": 1, "top_logprobs": []}
for tid in output_ids
]
}
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": generate_response,
"prompt_tokens": len(gen_req["token_ids"]),
"chat_request": {
"model": PARSER_MODEL,
"messages": messages,
"include_reasoning": include_reasoning,
},
},
)
assert resp.status_code == 200, resp.text
choice = resp.json()["choices"][0]
assert choice["message"]["content"] == "Five."
if include_reasoning:
assert choice["message"]["reasoning"]
assert len(choice["logprobs"]["content"]) == len(output_ids)
else:
assert choice["message"]["reasoning"] is None
assert choice["logprobs"] is None
@pytest.mark.asyncio
async def test_e2e_parsed_tool_call(parser_client, parser_tokenizer):
"""<tool_call> extracted into tool_calls field."""
@@ -1446,19 +1491,12 @@ _HERMES_TOOL_CALL_TEXT = (
)
@pytest.mark.parametrize(
("tool_choice", "text", "engine_finish_reason", "expected_finish_reason"),
[
("auto", _HERMES_TOOL_CALL_TEXT, "stop", "tool_calls"),
("auto", _HERMES_TOOL_CALL_TEXT, "length", "length"),
("required", "Hello there", "stop", "stop"),
],
)
def test_derender_chat_finish_reason_matches_chat_completion(
tool_choice, text, engine_finish_reason, expected_finish_reason
def _derender_hermes_chat(
text: str,
tool_choice: str = "auto",
engine_finish_reason: str = "stop",
parser_cls: type | None = None,
):
"""/derender applies the same finish_reason rule as the coupled chat
path: ``tool_calls`` only for a real call at a natural stop."""
tokenizer = get_tokenizer(MODEL_NAME)
renderer = MagicMock()
renderer.get_tokenizer.return_value = tokenizer
@@ -1474,6 +1512,8 @@ def test_derender_chat_finish_reason_matches_chat_completion(
enable_auto_tools=True,
tool_parser="hermes",
)
if parser_cls is not None:
derenderer.parser = parser_cls
request = ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "Weather in Paris?"}],
@@ -1488,6 +1528,54 @@ def test_derender_chat_finish_reason_matches_chat_completion(
)
(choice,) = derenderer._derender_chat(response, request)
return choice
@pytest.mark.parametrize(
("tool_choice", "text", "engine_finish_reason", "expected_finish_reason"),
[
("auto", _HERMES_TOOL_CALL_TEXT, "stop", "tool_calls"),
("auto", _HERMES_TOOL_CALL_TEXT, "length", "length"),
("required", "Hello there", "stop", "stop"),
],
)
def test_derender_chat_finish_reason_matches_chat_completion(
tool_choice, text, engine_finish_reason, expected_finish_reason
):
"""/derender applies the same finish_reason rule as the coupled chat
path: ``tool_calls`` only for a real call at a natural stop."""
choice = _derender_hermes_chat(text, tool_choice, engine_finish_reason)
assert choice.finish_reason == expected_finish_reason
assert bool(choice.message.tool_calls) == (text == _HERMES_TOOL_CALL_TEXT)
def test_derender_chat_tool_call_id_matches_chat_completion():
"""Batch /derender uses the coupled chat's default tool call ID format."""
choice = _derender_hermes_chat(_HERMES_TOOL_CALL_TEXT)
(tool_call,) = choice.message.tool_calls
assert tool_call.id.startswith("chatcmpl-tool-")
def test_derender_chat_keeps_parser_tool_call_id():
"""Batch /derender keeps the parser's tool call ID (Kimi K2 style
``functions.<name>:<idx>``), since chat templates reference it in
multi-turn history."""
class _KimiIdParser:
def __init__(self, *args, **kwargs):
pass
def parse(self, model_output, request, **kwargs):
call = FunctionCall(
id="functions.get_weather:1",
name="get_weather",
arguments='{"city": "Paris"}',
)
return None, None, [call]
choice = _derender_hermes_chat(_HERMES_TOOL_CALL_TEXT, parser_cls=_KimiIdParser)
(tool_call,) = choice.message.tool_calls
assert tool_call.id == "functions.get_weather:1"
tests/entrypoints/scale_out/derender/test_derender_parity.py
+176 -7
View File
@@ -142,12 +142,28 @@ def _tool_sig(response_choice: dict) -> list[tuple[str, dict]]:
]
def _tool_id_shapes(response_choice: dict) -> list[str | None]:
return [
_tool_call_id_shape(tc["id"])
for tc in (response_choice["message"].get("tool_calls") or [])
]
def _engine_finish_reason(coupled_finish_reason: str) -> str:
"""Undo the chat level `"tool_calls"` rewrite of `"stop"` so derender
gets what the engine reported and has to reapply the rewrite itself."""
if coupled_finish_reason == "tool_calls":
return "stop"
return coupled_finish_reason
def _assert_parity(coupled: dict, disagg: dict) -> None:
"""Both paths saw the same tokens, so they must agree unconditionally."""
c, d = coupled["choices"][0], disagg["choices"][0]
assert d["message"]["content"] == c["message"]["content"]
assert d["message"].get("reasoning") == c["message"].get("reasoning")
assert _tool_sig(d) == _tool_sig(c)
assert _tool_id_shapes(d) == _tool_id_shapes(c)
assert d["finish_reason"] == c["finish_reason"]
assert d["token_ids"] == c["token_ids"]
assert disagg["usage"]["prompt_tokens"] == coupled["usage"]["prompt_tokens"]
@@ -166,7 +182,7 @@ async def _run_parity_case(
client,
ch["token_ids"],
coupled["usage"]["prompt_tokens"],
ch["finish_reason"],
_engine_finish_reason(ch["finish_reason"]),
chat_request,
)
return coupled, disagg
@@ -316,13 +332,18 @@ async def test_parity_logprobs(client):
async def _coupled_stream(
client: httpx.AsyncClient, messages: list[dict], **extra
) -> tuple[list[dict], list[int]]:
client: httpx.AsyncClient,
messages: list[dict],
require_token_ids: bool = True,
**extra,
) -> tuple[list[dict], list[int] | None]:
"""Stream `/v1/chat/completions` and return the choices that carry
`token_ids` along with the prompt token IDs.
The leading role only chunk has no `token_ids` and no counterpart on
`/inference/v1/generate`, so it is dropped.
`/inference/v1/generate`, so it is dropped. With `require_token_ids`
off every choice is kept and `prompt_token_ids` may be `None`, for
requests where the coupled server hides token ids (hidden reasoning).
"""
choices: list[dict] = []
prompt_token_ids: list[int] | None = None
@@ -351,9 +372,12 @@ async def _coupled_stream(
if chunk.get("prompt_token_ids") is not None:
prompt_token_ids = chunk["prompt_token_ids"]
choices.extend(
ch for ch in chunk["choices"] if ch.get("token_ids") is not None
ch
for ch in chunk["choices"]
if not require_token_ids or ch.get("token_ids") is not None
)
assert prompt_token_ids is not None
if require_token_ids:
assert prompt_token_ids is not None
return choices, prompt_token_ids
@@ -390,7 +414,7 @@ async def _run_stream_parity_case(
{"model": MODEL, "messages": messages, "return_token_ids": True, **extra},
len(prompt_token_ids),
prompt_token_ids,
finish_reason="stop" if finish_reason == "tool_calls" else finish_reason,
finish_reason=_engine_finish_reason(finish_reason),
)
coupled = assemble_stream(coupled_choices)
for message in (coupled, disagg):
@@ -473,6 +497,151 @@ async def test_stream_parity_reasoning_and_tool_call(client, chunking):
pytest.skip("Model did not emit both a <think> block and a tool call")
# ---------------------------------------------------------------------------
# Streaming logprobs on the parser path
# ---------------------------------------------------------------------------
def _logprob_tokens(choice: dict) -> list[tuple[str, list[int] | None]] | None:
lp = choice.get("logprobs")
if lp is None or lp.get("content") is None:
return None
return [(e["token"], e.get("bytes")) for e in lp["content"]]
def _sampled_generate_logprobs(choice: dict) -> dict | None:
"""`GenerateLogProbs` for one coupled stream chunk from its `token_ids`
and per-token logprob values, without `top_logprobs`.
The coupled stream's `token_ids` are the integer ids a GPU-less worker
would have sent, so one stream is both the derender input and the string
reference. A second greedy run is not needed, and greedy decoding is not
bit-reproducible across calls on every platform anyway.
"""
lp = choice.get("logprobs")
if lp is None or lp.get("content") is None:
return None
content = lp["content"]
assert len(content) == len(choice["token_ids"]), choice
return {
"content": [
{"token_id": tid, "logprob": entry["logprob"], "top_logprobs": []}
for tid, entry in zip(choice["token_ids"], content)
]
}
async def _run_stream_logprob_parity_case(
client: httpx.AsyncClient, messages: list[dict], **extra
) -> tuple[dict, list[dict], list[dict]]:
"""Stream with logprobs through the coupled server, replay the same
chunks through the streaming derender endpoint and compare per chunk.
Returns the assembled coupled message, the coupled choices and the
derender choices.
"""
extra = {"logprobs": True, **extra}
coupled_choices, prompt_token_ids = await _coupled_stream(client, messages, **extra)
finish_reason = coupled_choices[-1]["finish_reason"]
derender_choices: list[dict] = []
disagg = await stream_chat_derender(
client,
[tid for ch in coupled_choices for tid in ch["token_ids"]],
[len(ch["token_ids"]) for ch in coupled_choices],
{"model": MODEL, "messages": messages, **extra},
len(prompt_token_ids),
prompt_token_ids,
finish_reason="stop" if finish_reason == "tool_calls" else finish_reason,
logprobs=[_sampled_generate_logprobs(ch) for ch in coupled_choices],
raw_choices=derender_choices,
)
coupled = assemble_stream(coupled_choices)
for message in (coupled, disagg):
for tc in message["tool_calls"]:
tc["id"] = _tool_call_id_shape(tc["id"])
assert disagg == coupled
assert len(derender_choices) == len(coupled_choices)
for ref, der in zip(coupled_choices, derender_choices):
assert _logprob_tokens(der) == _logprob_tokens(ref), (ref, der)
return coupled, coupled_choices, derender_choices
@pytest.mark.asyncio
async def test_stream_parity_logprobs_reasoning(client):
"""Per-chunk logprobs on the parser path match the coupled stream across
the reasoning -> content transition."""
messages = [{"role": "user", "content": "What is 17 times 23? Think it through."}]
coupled, _, derender_choices = await _run_stream_logprob_parity_case(
client, messages, include_reasoning=True, max_tokens=256
)
if not coupled["reasoning"]:
pytest.skip("Model did not emit a <think> block")
assert any(
ch["delta"].get("reasoning") and ch.get("logprobs") for ch in derender_choices
)
@pytest.mark.asyncio
async def test_stream_parity_logprobs_tool_call(client):
"""Chunks the parser swallows (tool call markup) still carry their
logprobs, on both sides, so no token's logprobs are lost."""
messages = [{"role": "user", "content": "What's the weather in Paris?"}]
coupled, _, derender_choices = await _run_stream_logprob_parity_case(
client,
messages,
tools=TOOLS,
tool_choice=FORCE_WEATHER_TOOL,
include_reasoning=True,
max_tokens=512,
)
if not coupled["tool_calls"]:
pytest.skip("Model did not emit a tool call")
swallowed = [
ch
for ch in derender_choices
if ch.get("logprobs")
and not ch["delta"].get("content")
and not ch["delta"].get("reasoning")
and not ch["delta"].get("tool_calls")
]
assert swallowed, "expected at least one buffered-markup chunk with logprobs"
@pytest.mark.asyncio
async def test_stream_parity_logprobs_hidden_reasoning(client):
"""With `include_reasoning=False` and a parser configured, neither side
returns logprobs on any chunk; derender still emits the content."""
messages = [{"role": "user", "content": "What is 17 times 23? Think it through."}]
extra = {"logprobs": True, "max_tokens": 256}
# Token ids come from a shown-reasoning run: the coupled server hides
# them along with the logprobs when reasoning is hidden.
shown_choices, prompt_token_ids = await _coupled_stream(
client, messages, include_reasoning=True, **extra
)
hidden_choices, _ = await _coupled_stream(
client, messages, require_token_ids=False, include_reasoning=False, **extra
)
assert all(ch.get("logprobs") is None for ch in hidden_choices)
assert assemble_stream(hidden_choices)["reasoning"] is None
finish_reason = shown_choices[-1]["finish_reason"]
derender_choices: list[dict] = []
disagg = await stream_chat_derender(
client,
[tid for ch in shown_choices for tid in ch["token_ids"]],
[len(ch["token_ids"]) for ch in shown_choices],
{"model": MODEL, "messages": messages, "include_reasoning": False, **extra},
len(prompt_token_ids),
prompt_token_ids,
finish_reason="stop" if finish_reason == "tool_calls" else finish_reason,
logprobs=[_sampled_generate_logprobs(ch) for ch in shown_choices],
raw_choices=derender_choices,
)
assert all(ch.get("logprobs") is None for ch in derender_choices)
shown = assemble_stream(shown_choices)
assert disagg["reasoning"] is None
assert disagg["content"] == shown["content"]
# ---------------------------------------------------------------------------
# Inline text level: generate(output_mode="text") == derender(generate(...))
# ---------------------------------------------------------------------------
tests/entrypoints/scale_out/derender/test_derender_stream.py
+153
View File
@@ -1427,6 +1427,159 @@ class TestDerenderChatStreamParsed:
assert delta.reasoning is None
assert delta.content == "c"
@pytest.mark.asyncio
async def test_logprobs_attached_to_each_parsed_delta(
self, parsed_derenderer, tokenizer
):
"""A chunk's logprobs ride on whatever delta the parser emits for it
(reasoning, content or tool call), with token ids decoded, as on
`/v1/chat/completions`."""
chunks = [
[_FakeParser.REASON],
[_FakeParser.CONTENT],
[_FakeParser.TOOL_START, _FakeParser.TOOL_ARG],
]
state = None
seen = []
for ids in chunks:
chunk, state = await parsed_derenderer.derender_chat_stream(
model=MODEL_NAME,
generate_chunk=_make_stream_chunk(
ids, logprobs=_generate_logprobs(ids)
),
state=state,
chat_request=_chat_request(),
prompt_token_ids=[1, 2],
)
choice = chunk.choices[0]
assert choice.logprobs is not None and choice.logprobs.content is not None
assert [e.token for e in choice.logprobs.content] == (
convert_ids_list_to_tokens(tokenizer, ids)
)
seen.append(choice.delta)
assert seen[0].reasoning == "r"
assert seen[1].content == "c"
assert (
seen[2].tool_calls and seen[2].tool_calls[0].function.name == "get_weather"
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("request_kwargs", "expected_candidates"),
[({"logprobs": True, "top_logprobs": 2}, 2), ({}, 4)],
)
async def test_top_logprobs_follow_the_chat_cut(
self, parsed_derenderer, request_kwargs, expected_candidates
):
"""The parsed path trims generate's candidate list to the request's
`top_logprobs` like the plain path does, and keeps every candidate
when the request does not ask for logprobs."""
ids = [_FakeParser.CONTENT]
logprobs = _generate_logprobs(ids)
logprobs["content"][0]["top_logprobs"] = [
{"token_id": tid, "logprob": -float(rank)}
for rank, tid in enumerate([ids[0], 11, 12, 13])
]
chunk, _ = await parsed_derenderer.derender_chat_stream(
model=MODEL_NAME,
generate_chunk=_make_stream_chunk(ids, logprobs=logprobs),
chat_request=_chat_request(**request_kwargs),
prompt_token_ids=[1, 2],
)
content = chunk.choices[0].logprobs.content
assert len(content) == 1
assert len(content[0].top_logprobs) == expected_candidates
@pytest.mark.asyncio
async def test_logprobs_kept_when_parser_swallows_chunk(
self, parsed_derenderer, tokenizer
):
"""A chunk the parser emits nothing for (buffered markup) still
carries its logprobs on an empty delta, so no token's logprobs are
lost. Mirrors the coupled endpoint, which emits a chunk whenever
logprobs are present even if the delta is empty."""
swallowed = [_FakeParser.CONTENT + 1] # unknown to _FakeParser -> None
chunk, _ = await parsed_derenderer.derender_chat_stream(
model=MODEL_NAME,
generate_chunk=_make_stream_chunk(
swallowed, logprobs=_generate_logprobs(swallowed)
),
chat_request=_chat_request(),
prompt_token_ids=[1, 2],
)
choice = chunk.choices[0]
assert choice.delta.content is None and not choice.delta.tool_calls
assert [e.token for e in choice.logprobs.content] == (
convert_ids_list_to_tokens(tokenizer, swallowed)
)
@pytest.mark.asyncio
async def test_include_reasoning_false_drops_logprobs_on_every_chunk(
self, parsed_derenderer
):
"""Hidden reasoning suppresses logprobs on reasoning *and* content
chunks (decoded logprob tokens would leak the reasoning), the same
rule `/v1/chat/completions` applies with a parser configured."""
state = None
for ids in ([_FakeParser.REASON], [_FakeParser.CONTENT]):
chunk, state = await parsed_derenderer.derender_chat_stream(
model=MODEL_NAME,
generate_chunk=_make_stream_chunk(
ids, logprobs=_generate_logprobs(ids)
),
state=state,
chat_request=_chat_request(include_reasoning=False),
prompt_token_ids=[1, 2],
)
assert chunk.choices[0].logprobs is None
assert chunk.choices[0].delta.content == "c"
@pytest.mark.asyncio
async def test_logprobs_multibyte_across_chunks(self, parsed_derenderer, tokenizer):
"""Byte-fallback correction crosses chunk boundaries on the parser
path too, via the carried `logprob_context_token_ids`, which starts
empty like the engine's and never includes the prompt."""
token_ids = tokenizer.encode("👍", add_special_tokens=False)
if len(token_ids) < 2 or "�" not in tokenizer.decode([token_ids[-1]]):
pytest.skip("Tokenizer does not byte-split this character")
state = None
last_chunk = None
for tid in token_ids:
last_chunk, state = await parsed_derenderer.derender_chat_stream(
model=MODEL_NAME,
generate_chunk=_make_stream_chunk(
[tid], logprobs=_generate_logprobs([tid])
),
state=state,
chat_request=_chat_request(),
prompt_token_ids=[1, 2, 3, 4, 5],
)
assert last_chunk is not None and state is not None
final_entry = last_chunk.choices[0].logprobs.content[0]
assert not final_entry.token.endswith("�"), final_entry.token
assert state.logprob_context_token_ids == token_ids[-4:]
@pytest.mark.asyncio
async def test_logprob_context_advances_without_logprobs(self, parsed_derenderer):
"""The context tail tracks sampled ids whether or not the chunk
carries logprobs, so a later chunk that does can still be corrected."""
ids = [_FakeParser.REASON, _FakeParser.CONTENT, _FakeParser.CONTENT]
_, state = await parsed_derenderer.derender_chat_stream(
model=MODEL_NAME,
generate_chunk=_make_stream_chunk(ids),
chat_request=_chat_request(),
prompt_token_ids=[1, 2],
)
assert state.logprob_context_token_ids == ids
_, state = await parsed_derenderer.derender_chat_stream(
model=MODEL_NAME,
generate_chunk=_make_stream_chunk([], finish_reason="stop"),
state=state,
chat_request=_chat_request(),
prompt_token_ids=[1, 2],
)
assert state.logprob_context_token_ids == ids
# ---------------------------------------------------------------------------
# Harmony / GPT-OSS replay — unit, no server
tests/entrypoints/scale_out/derender/utils.py
+11
View File
@@ -52,6 +52,8 @@ async def stream_chat_derender(
prompt_token_ids: list[int],
finish_reason: str = "stop",
on_chunk: Callable[[list[dict]], None] | None = None,
logprobs: list[dict | None] | None = None,
raw_choices: list[dict] | None = None,
) -> dict:
"""Feed `output_ids` through the streaming chat derender endpoint in
the given `chunk_sizes`, threading `stream_state` across calls and
@@ -64,7 +66,13 @@ async def stream_chat_derender(
`tool_calls` assembled so far, letting callers assert properties of the
intermediate deltas (e.g. monotonic argument growth) rather than only
the final assembled result.
`logprobs` supplies one `GenerateLogProbs` payload (or `None`) per chunk
and `raw_choices`, if given, collects every derendered choice so callers
can compare per-chunk fields such as `logprobs`.
"""
if logprobs is not None:
assert len(logprobs) == len(chunk_sizes)
state = None
choices: list[dict] = []
pos = 0
@@ -84,6 +92,7 @@ async def stream_chat_derender(
"index": 0,
"token_ids": tids,
"finish_reason": finish_reason if is_last else None,
"logprobs": logprobs[i] if logprobs is not None else None,
}
],
},
@@ -101,6 +110,8 @@ async def stream_chat_derender(
assert data["chunk"]["choices"][0]["token_ids"] == tids
state = data["stream_state"]
choices.extend(data["chunk"]["choices"])
if raw_choices is not None:
raw_choices.extend(data["chunk"]["choices"])
if on_chunk is not None:
on_chunk(assemble_stream(choices)["tool_calls"])
tests/evals/gsm8k/configs/DeepSeek-R1-DP_MI325.yaml
-12
View File
@@ -1,12 +0,0 @@
model_name: "deepseek-ai/DeepSeek-R1"
accuracy_threshold: 0.95
num_questions: 1319
num_fewshot: 5
startup_max_wait_seconds: 1200
server_args: >-
--enforce-eager
--max-model-len 4096
--data-parallel-size 8
--enable-expert-parallel
--attention-backend=TRITON_ATTN
--speculative-config '{"method":"mtp","num_speculative_tokens":3}'
tests/evals/gsm8k/configs/DeepSeek-R1-TP_MI325.yaml
-12
View File
@@ -1,12 +0,0 @@
model_name: "deepseek-ai/DeepSeek-R1"
accuracy_threshold: 0.95
num_questions: 1319
num_fewshot: 5
startup_max_wait_seconds: 1200
server_args: >-
--enforce-eager
--max-model-len 4096
--tensor-parallel-size 8
--enable-expert-parallel
--attention-backend=TRITON_ATTN
--speculative-config '{"method":"mtp","num_speculative_tokens":3}'
tests/evals/gsm8k/configs/DeepSeek-V3.2-DP_MI325.yaml
-12
View File
@@ -1,12 +0,0 @@
model_name: "deepseek-ai/DeepSeek-V3.2"
accuracy_threshold: 0.95
num_questions: 1319
num_fewshot: 5
startup_max_wait_seconds: 1200
server_args: >-
--enforce-eager
--max-model-len 4096
--data-parallel-size 8
--enable-expert-parallel
--attention-backend=TRITON_ATTN
--speculative-config '{"method":"mtp","num_speculative_tokens":3}'
tests/evals/gsm8k/configs/DeepSeek-V3.2-TP_MI325.yaml
-12
View File
@@ -1,12 +0,0 @@
model_name: "deepseek-ai/DeepSeek-V3.2"
accuracy_threshold: 0.95
num_questions: 1319
num_fewshot: 5
startup_max_wait_seconds: 1200
server_args: >-
--enforce-eager
--max-model-len 4096
--tensor-parallel-size 8
--enable-expert-parallel
--attention-backend=TRITON_ATTN
--speculative-config '{"method":"mtp","num_speculative_tokens":3}'
tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16_MI355.yaml
+1
View File
@@ -12,4 +12,5 @@ server_args: >-
--enable-expert-parallel
--mamba-backend triton
--moe-backend triton
--attention-backend ROCM_AITER_FA
--speculative-config '{"method":"mtp","num_speculative_tokens":5}'
tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-FP8_MI355.yaml
+1
View File
@@ -12,4 +12,5 @@ server_args: >-
--enable-expert-parallel
--mamba-backend triton
--moe-backend triton
--attention-backend ROCM_AITER_FA
--speculative-config '{"method":"mtp","num_speculative_tokens":5}'
tests/evals/gsm8k/configs/models-mi3xx.txt
-8
View File
@@ -1,8 +0,0 @@
DeepSeek-R1-TP_MI325.yaml
DeepSeek-R1-DP_MI325.yaml
DeepSeek-V3.2-TP_MI325.yaml
DeepSeek-V3.2-DP_MI325.yaml
Qwen3-30B-A3B-NVFP4.yaml
Qwen3-30B-A3B-NVFP4-quark.yaml
Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml
Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml
tests/kernels/core/test_activation.py
-40
View File
@@ -353,46 +353,6 @@ def test_silu_and_mul_with_clamp(
)
@pytest.mark.parametrize("linear_beta", [-1.0, 2.0])
@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16])
@torch.inference_mode()
def test_masked_situ_and_mul(
default_vllm_config,
linear_beta: float,
dtype: torch.dtype,
) -> None:
"""Masked SITU computes valid expert rows and preserves padded zeros."""
device = CUDA_DEVICES[0]
num_experts, max_num_tokens, d = 4, 7, 512
beta = 1.5
input = torch.randn(num_experts, max_num_tokens, 2 * d, dtype=dtype, device=device)
expert_num_tokens = torch.tensor([0, 1, 4, 7], dtype=torch.int32, device=device)
output = torch.zeros(num_experts, max_num_tokens, d, dtype=dtype, device=device)
torch.ops._C.masked_situ_and_mul(
output, input, expert_num_tokens, beta, linear_beta
)
gate, up = input.float().chunk(2, dim=-1)
expected = beta * torch.tanh(gate / beta) * torch.sigmoid(gate)
if linear_beta > 0:
up = linear_beta * torch.tanh(up / linear_beta)
expected = (expected * up).to(dtype)
for expert, num_tokens in enumerate(expert_num_tokens.cpu().tolist()):
torch.testing.assert_close(
output[expert, :num_tokens],
expected[expert, :num_tokens],
atol=get_default_atol(output),
rtol=get_default_rtol(output),
)
assert torch.count_nonzero(output[expert, num_tokens:]) == 0
opcheck(
torch.ops._C.masked_situ_and_mul,
(output, input, expert_num_tokens, beta, linear_beta),
)
MOE_ACTIVATION_CASES = [
pytest.param(MoEActivation.SILU, ApplyMoEActivationConfig(), id="silu"),
pytest.param(
tests/kernels/mamba/test_causal_conv1d.py
+112
View File
@@ -2,6 +2,8 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterator
import pytest
import torch
import torch.nn.functional as F
@@ -358,3 +360,113 @@ def test_causal_conv1d_varlen(
)
unpadded_out = out[:, : out_ref_tensor.shape[-1]]
assert torch.allclose(unpadded_out, out_ref_tensor, rtol=rtol, atol=atol)
@pytest.mark.parametrize("itype", [torch.bfloat16])
@pytest.mark.parametrize("silu_activation", [True])
@pytest.mark.parametrize("has_bias", [True])
@pytest.mark.parametrize("width", [4])
@pytest.mark.parametrize("dim", [64, 2048])
@pytest.mark.parametrize("q_lens_cfg", [[3, 1, 0, 4, 2], [1, 5, 5, 1], [2, 0, 0, 3]])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason="CPU kernel for causal_conv1d_update ignores num_accepted_tokens.",
)
def test_causal_conv1d_update_ragged_varlen(
q_lens_cfg, dim, width, has_bias, silu_activation, itype
):
mk_inputs = lambda: _create_ragged_varlen_inputs(
q_lens_cfg, dim, width, has_bias, silu_activation, itype
)
inputs, state = mk_inputs()
actual = causal_conv1d_update(conv_state=state, **inputs)
ref_inputs, ref_state = mk_inputs()
ref_outs = _create_ragged_varlen_ref_outs(conv_state=ref_state, **ref_inputs)
for i, ref_out in enumerate(ref_outs):
start = int(inputs["query_start_loc"][i])
end = int(inputs["query_start_loc"][i + 1])
entry = inputs["conv_state_indices"][i : i + 1]
assert torch.allclose(actual[start:end], ref_out, rtol=1e-2, atol=5e-2)
assert torch.equal(state[entry], ref_state[entry])
def _create_ragged_varlen_inputs(
q_lens_cfg, dim, width, has_bias, silu_activation, itype
):
device = DEVICE
set_random_seed(0)
batch_size = len(q_lens_cfg)
max_query_len = 5
assert max(q_lens_cfg) <= max_query_len
state_len = width - 1 + max_query_len
total_entries = 10 * batch_size
q_lens = torch.tensor(q_lens_cfg, dtype=torch.int32, device=device)
query_start_loc = torch.cat(
[
torch.zeros(1, dtype=torch.int32, device=device),
torch.cumsum(q_lens, dim=0, dtype=torch.int32),
]
)
total_tokens = int(q_lens.sum())
x = torch.randn(total_tokens, dim, device=device, dtype=itype)
conv_state = torch.randn(
total_entries, state_len, dim, device=device, dtype=itype
).transpose(1, 2)
weight = torch.randn(dim, width, device=device, dtype=itype)
bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
activation = None if not silu_activation else "silu"
conv_state_indices = (
torch.randperm(total_entries - 1, device=device)[:batch_size] + 1
).to(torch.int32)
num_accepted_tokens = torch.randint(
1, max_query_len + 1, (batch_size,), dtype=torch.int32, device=device
)
inputs = dict(
x=x,
weight=weight,
bias=bias,
activation=activation,
conv_state_indices=conv_state_indices,
num_accepted_tokens=num_accepted_tokens,
query_start_loc=query_start_loc,
max_query_len=max_query_len,
)
return inputs, conv_state
def _create_ragged_varlen_ref_outs(
x,
conv_state,
weight,
bias,
activation,
conv_state_indices,
num_accepted_tokens,
query_start_loc,
max_query_len,
) -> Iterator[torch.Tensor]:
for seq_idx in range(len(conv_state_indices)):
start = int(query_start_loc[seq_idx])
end = int(query_start_loc[seq_idx + 1])
entry = conv_state_indices[seq_idx : seq_idx + 1]
if start == end:
yield x[start:end]
continue
out_ref = causal_conv1d_update(
x[start:end],
conv_state,
weight,
bias,
activation=activation,
conv_state_indices=entry,
num_accepted_tokens=num_accepted_tokens[seq_idx : seq_idx + 1],
query_start_loc=query_start_loc.new_tensor([0, end - start]),
max_query_len=max_query_len,
)
yield out_ref
tests/kernels/mamba/test_mamba_ssm_ssd.py
+47
View File
@@ -6,6 +6,10 @@ import torch
import torch.nn.functional as F
from einops import rearrange, repeat
from vllm.model_executor.layers.mamba.ops.ssd_chunk_scan import (
_chunk_scan_fwd,
_chunk_scan_fwd_kernel,
)
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
mamba_chunk_scan_combined_varlen,
)
@@ -578,3 +582,46 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
rtol=rtol,
msg=lambda x, i=i: f"seq{i} state " + x,
)
def test_chunk_scan_fwd_every_config_matches(monkeypatch):
"""The autotuner picks a config by speed only, so every config in the list
must give the same finite output. On sm90, configs that pipeline a wgmma
with `cb` in registers wrote NaN into whole rows (#61603)."""
from triton.runtime.errors import OutOfResources
set_random_seed(0)
nheads, headdim, ngroups, dstate, chunk_size = 64, 64, 8, 128, 128
cu_chunk_seqlens = torch.tensor([0, 128, 181], dtype=torch.int32, device=DEVICE)
seqlen, nchunks = 181, 2
x = torch.randn(seqlen, nheads, headdim, dtype=torch.bfloat16, device=DEVICE)
C = torch.randn(seqlen, ngroups, dstate, dtype=torch.bfloat16, device=DEVICE)
cb = torch.randn(nchunks, ngroups, chunk_size, chunk_size, device=DEVICE)
dt = torch.rand(nheads, nchunks, chunk_size, device=DEVICE) * 0.1
dA_cumsum = -torch.rand(nheads, nchunks, chunk_size, device=DEVICE).cumsum(-1)
states = torch.randn(nchunks, nheads, headdim, dstate, device=DEVICE) * 0.1
seq_idx = torch.zeros(nchunks, dtype=torch.int32, device=DEVICE)
D = torch.randn(nheads, headdim, device=DEVICE) * 0.1
ref = None
for config in list(_chunk_scan_fwd_kernel.configs):
monkeypatch.setattr(_chunk_scan_fwd_kernel, "configs", [config])
for _ in range(10):
out = torch.empty_like(x)
try:
_chunk_scan_fwd(
cb, x, dt, dA_cumsum, C, states, cu_chunk_seqlens, out, seq_idx, D
)
except OutOfResources:
# The autotuner skips configs that do not fit this device.
break
assert torch.isfinite(out).all(), f"non-finite output with {config}"
if ref is None:
ref = out
torch.testing.assert_close(
out,
ref,
atol=2e-2,
rtol=2e-2,
msg=lambda m, config=config: f"{config}: {m}",
)
tests/kernels/mhc/test_mhc_kernels.py
+13 -1
View File
@@ -10,7 +10,6 @@ import torch.nn.functional as F
import vllm.model_executor.kernels.mhc # noqa: F401
import vllm.model_executor.layers.mhc as mhc_layers
from vllm.model_executor.kernels.mhc.tilelang import (
_torch_hc_prenorm_gemm,
mhc_post_tilelang,
mhc_pre_delayed_tilelang,
)
@@ -1002,6 +1001,19 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult, fused_norm):
torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2)
def _torch_hc_prenorm_gemm(
x: torch.Tensor,
fn: torch.Tensor,
out: torch.Tensor,
sqrsum: torch.Tensor,
) -> None:
assert out.shape[0] == 1
assert sqrsum.shape[0] == 1
x_float = x.float()
out[0].copy_(x_float @ fn.t())
sqrsum[0].copy_(x_float.square().sum(dim=-1))
@pytest.mark.skipif(
not HAS_TILELANG_MHC,
reason="TileLang MHC support required",
tests/kernels/moe/modular_kernel_tools/parallel_utils.py
+16 -2
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import dataclasses
import os
import time
import traceback
from collections.abc import Callable
from typing import Any, Concatenate
@@ -23,6 +24,8 @@ from vllm.utils.network_utils import get_open_port
P = ParamSpec("P")
LAUNCH_TIMEOUT_S = 900
@dataclasses.dataclass
class ProcessGroupInfo:
@@ -148,7 +151,7 @@ def parallel_launch_with_config(
*args: P.args,
**kwargs: P.kwargs,
) -> None:
spawn(
context = spawn(
_worker_parallel_launch,
args=(
world_size,
@@ -162,5 +165,16 @@ def parallel_launch_with_config(
)
+ args,
nprocs=world_size,
join=True,
join=False,
)
# A worker stuck in communicator setup never exits, so bound the wait.
deadline = time.monotonic() + LAUNCH_TIMEOUT_S
while not context.join(timeout=max(0.0, deadline - time.monotonic())):
if time.monotonic() >= deadline:
for process in context.processes:
process.kill()
for process in context.processes:
process.join()
raise TimeoutError(
f"Parallel workers did not finish within {LAUNCH_TIMEOUT_S}s"
)
tests/kernels/moe/test_moe_layer.py
+11 -1
View File
@@ -127,6 +127,14 @@ if has_nixl_ep():
BACKENDS += ["nixl_ep"]
DEEPEP_BACKENDS = {"deepep_high_throughput", "deepep_low_latency"}
# These tests run on one node, where NVSHMEM reaches peers over NVLink, so
# keep it off the NICs to avoid depending on IB setup on shared CI hosts.
DEEPEP_SINGLE_NODE_NVSHMEM_ENV = {
"NVSHMEM_REMOTE_TRANSPORT": "none",
# DeepEP forces IBGDA on, so also hide the NICs from it.
"NVSHMEM_HCA_LIST": "none",
}
MORI_BACKENDS = {"mori_high_throughput", "mori_low_latency"}
QUANT_METHODS = [
@@ -2187,6 +2195,8 @@ def test_moe_layer(
rocm_deepep = current_platform.is_rocm() and backend in DEEPEP_BACKENDS
parallel_worker = _parallel_worker_rocm_deepep if rocm_deepep else _parallel_worker
cuda_deepep = current_platform.is_cuda() and backend in DEEPEP_BACKENDS
env_dict = DEEPEP_SINGLE_NODE_NVSHMEM_ENV if cuda_deepep else None
launch_failures: list[str] = []
try:
@@ -2197,7 +2207,7 @@ def test_moe_layer(
world_size,
parallel_worker,
vllm_config,
None,
env_dict,
test_config_batch,
verbosity,
failure_report_path=failure_report_path,
tests/kernels/moe/test_routing.py
+14
View File
@@ -582,6 +582,20 @@ def test_fused_topk_bias(
assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
@pytest.mark.parametrize("scoring_func", [0, 1])
def test_fused_grouped_topk_accepts_idle_rank(scoring_func):
"""Zero-token routing must leave CUDA ready for the next batch."""
if not current_platform.is_cuda():
pytest.skip("CUDA required")
from vllm import _custom_ops as ops
scores = torch.empty((0, 8), device="cuda", dtype=torch.float32)
bias = torch.zeros(8, device="cuda")
weights, ids = ops.grouped_topk(scores, 2, 1, 2, True, 1.0, bias, scoring_func)
torch.testing.assert_close(weights, scores.new_empty((0, 2)))
torch.testing.assert_close(ids, scores.new_empty((0, 2), dtype=torch.int32))
@pytest.mark.parametrize("m,k", MK_S)
@pytest.mark.parametrize("top_k", TOP_KS)
@pytest.mark.parametrize(
tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py
+6 -2
View File
@@ -116,7 +116,7 @@ def test_silu_mul_fp8_quant_deep_gemm(T: int, N: int):
@pytest.mark.parametrize("scale_format", ["packed", "float32", "ceil_ue8m0"])
@pytest.mark.parametrize("hidden_size", [256, 2048])
@pytest.mark.parametrize("clamp_limit", [None, 10.0])
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA graph test")
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA/ROCm graph test")
def test_silu_quant_skips_expert_padding_on_graph_replay(
counts: list[int],
group_size: int,
@@ -135,7 +135,10 @@ def test_silu_quant_skips_expert_padding_on_graph_replay(
x *= 16
clean_x = x.clone()
ends = torch.empty(len(counts), device=DEVICE, dtype=torch.int32)
output = torch.empty((capacity, hidden_size), device=DEVICE, dtype=FLOAT8_DTYPE)
output_dtype = (
FLOAT8_DTYPE if scale_format == "packed" else current_platform.fp8_dtype()
)
output = torch.empty((capacity, hidden_size), device=DEVICE, dtype=output_dtype)
def set_counts(values):
positions = []
@@ -187,6 +190,7 @@ def test_silu_quant_skips_expert_padding_on_graph_replay(
output.view(torch.uint8).fill_(91)
graph.replay()
reference_q, reference_s = quantize(x)
assert actual_q.dtype == reference_q.dtype
torch.testing.assert_close(
actual_q.view(torch.uint8)[live],
reference_q.view(torch.uint8)[live],
tests/kernels/test_fp32_router_gemm.py
+57 -15
View File
@@ -3,11 +3,11 @@
"""Tests for fp32_router_gemm kernel: activation×weight→fp32.
Supported (hidden_size, num_experts) pairs:
(3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
(3072, 256), (6144, 128), (5120, 256), (4096, 352), (8192, 256)
Correctness baseline: F.linear in float32. Every M in [1, 32] is covered so
Correctness baseline: F.linear in float32. Every supported M is covered so
all tuned geometries (wide-block, experts-per-block, token-group; boundaries
at M=4/5, odd/even, M=15/16) are exercised on Blackwell, and the legacy
at M=4/5, M=9/10, M=12/13, odd/even, M=15/16) are exercised on Blackwell, and the legacy
128/1 geometry everywhere else.
"""
@@ -18,8 +18,12 @@ from vllm._custom_ops import fp32_router_gemm
from vllm.platforms import current_platform
# (hidden_size, num_experts)
SHAPES = [(3072, 256), (6144, 128)]
ALL_M = list(range(1, 33))
SHAPES = [(3072, 256), (6144, 128), (5120, 256), (4096, 352), (8192, 256)]
CASES = [
(m, hidden, experts)
for hidden, experts in SHAPES
for m in range(1, 49 if hidden == 5120 else 33)
]
# Absolute tolerance for fp32 kernel vs float64 reference
ATOL_FP32 = 2e-4
ATOL_BF16 = 2e-2 # bf16 activation has lower precision
@@ -40,8 +44,7 @@ def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
return torch.nn.functional.linear(mat_a.float(), mat_b.float())
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
@pytest.mark.parametrize("num_tokens", ALL_M)
@pytest.mark.parametrize("num_tokens,hidden_dim,num_experts", CASES)
def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
"""fp32 activation → fp32 output should match reference closely."""
_requires_sm90()
@@ -58,8 +61,7 @@ def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
torch.testing.assert_close(out, ref, atol=ATOL_FP32, rtol=0)
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
@pytest.mark.parametrize("num_tokens", ALL_M)
@pytest.mark.parametrize("num_tokens,hidden_dim,num_experts", CASES)
def test_bf16_activation(num_tokens: int, hidden_dim: int, num_experts: int):
"""bf16 activation → fp32 output should match reference within bf16 error."""
_requires_sm90()
@@ -124,11 +126,11 @@ def test_topk_routing_consistency(num_tokens: int, hidden_dim: int, num_experts:
)
def test_zero_tokens_returns_empty():
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
def test_zero_tokens_returns_empty(hidden_dim, num_experts):
"""M=0 is a graceful no-op returning an empty [0, E] tensor."""
_requires_sm90()
device = torch.device("cuda")
hidden_dim, num_experts = SHAPES[0]
mat_a = torch.empty(0, hidden_dim, dtype=torch.float32, device=device)
mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
out = fp32_router_gemm(mat_a, mat_b)
@@ -136,17 +138,19 @@ def test_zero_tokens_returns_empty():
assert out.dtype == torch.float32
def test_rejects_invalid_inputs():
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
def test_rejects_invalid_inputs(hidden_dim, num_experts):
"""The entry must fail loudly, never compute silently wrong results."""
_requires_sm90()
device = torch.device("cuda")
hidden_dim, num_experts = SHAPES[0]
mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
# num_tokens > 32 (beyond the instantiated range)
# Beyond the instantiated range for this shape.
invalid_tokens = 49 if hidden_dim == 5120 else 33
with pytest.raises(Exception, match="num_tokens"):
fp32_router_gemm(
torch.randn(33, hidden_dim, dtype=torch.float32, device=device), mat_b
torch.randn(invalid_tokens, hidden_dim, dtype=torch.float32, device=device),
mat_b,
)
# unsupported (hidden_dim, num_experts) pair
@@ -173,3 +177,41 @@ def test_rejects_invalid_inputs():
fp32_router_gemm(
torch.randn(4, hidden_dim, dtype=torch.float16, device=device), mat_b
)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize(
"tokens", [*range(1, 17), 17, 24, 32, 33, 40, 48, 56, 64, 65, 128]
)
@pytest.mark.parametrize(
"hidden_dim,num_experts", [(5120, 256), (4096, 352), (8192, 256)]
)
@torch.inference_mode()
def test_router_retains_fp32_weights_and_logits(
default_vllm_config, dist_init, tokens, hidden_dim, num_experts
):
"""Low-latency routing must retain FP32 weight precision and logits."""
from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
from vllm.platforms import current_platform
if not current_platform.is_device_capability_family(100):
pytest.skip("SM100 required")
torch.manual_seed(42)
router = GateLinear(
hidden_dim, num_experts, params_dtype=torch.float32, out_dtype=torch.float32
).cuda()
router.weight.normal_(std=0.01)
x = torch.randn(tokens, hidden_dim, device="cuda", dtype=torch.bfloat16)
expected = (x.double() @ router.weight.double().T).float()
actual, bias = router(x)
assert bias is None and actual.dtype == torch.float32
torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-5)
torch.testing.assert_close(actual.topk(8).indices, expected.topk(8).indices)
if tokens in (1, 16, 32, 33, 48, 56):
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
captured, _ = router(x)
x.normal_()
graph.replay()
expected = (x.double() @ router.weight.double().T).float()
torch.testing.assert_close(captured, expected, rtol=2e-5, atol=2e-5)
tests/kernels/test_ll_bf16_gemm.py
+25 -1
View File
@@ -566,7 +566,7 @@ def test_invalid_output_dtype():
with pytest.raises(ValueError, match="output_dtype=torch.float32"):
from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm
ll_bf16_gemm(a, b, output_dtype=torch.bfloat16)
ll_bf16_gemm(a, b, output_dtype=torch.float16)
def test_cache_miss_compiles_dotprod():
@@ -581,5 +581,29 @@ def test_cache_miss_compiles_dotprod():
_assert_close(out, _ref(a, b), context="cache miss dotprod")
@pytest.mark.parametrize("M", [1, 2, 8, 16])
@pytest.mark.parametrize("N", [1792, 3584])
@torch.inference_mode()
def test_bf16_output_and_fp32_cache_coexist(M, N):
"""BF16 epilogues retain FP32 accumulation and can replay with new inputs."""
from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm
torch.manual_seed(42)
x = torch.randn(M, 5120, device="cuda", dtype=torch.bfloat16)
w = torch.randn(N, 5120, device="cuda", dtype=torch.bfloat16) * 0.01
for dtype in (torch.float32, torch.bfloat16):
actual = ll_bf16_gemm(x, w, output_dtype=dtype)
assert actual.dtype == dtype
reference = (x.double() @ w.double().T).to(dtype)
torch.testing.assert_close(actual, reference, rtol=0.008, atol=1e-5)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
actual = ll_bf16_gemm(x, w, output_dtype=torch.bfloat16)
x.normal_()
graph.replay()
reference = (x.double() @ w.double().T).bfloat16()
torch.testing.assert_close(actual, reference, rtol=0.008, atol=1e-5)
if __name__ == "__main__":
pytest.main([__file__, "-v"])
tests/model_executor/test_flashinfer_autotune_warmup.py
+267 -4
View File
@@ -4,7 +4,7 @@
import json
import sys
from collections import defaultdict
from contextlib import nullcontext
from contextlib import contextmanager, nullcontext
from inspect import signature
from pathlib import Path
from types import ModuleType, SimpleNamespace
@@ -12,8 +12,10 @@ from typing import Any
from unittest.mock import Mock, call, patch
import pytest
import torch
from vllm.model_executor.warmup import kernel_warmup as warmup
import vllm.model_executor.warmup.kernel_warmup as warmup
import vllm.utils.flashinfer as fi_utils
from vllm.model_executor.warmup.kernel_warmup import (
_flashinfer_autotune_token_counts,
_run_flashinfer_autotune_dummy_runs,
@@ -45,7 +47,12 @@ def _make_runner(modules, *, max_tokens: int = 8192):
scheduler_config=SimpleNamespace(max_num_batched_tokens=max_tokens),
vllm_config=SimpleNamespace(
attention_config=SimpleNamespace(hisparse_config=None),
parallel_config=SimpleNamespace(data_parallel_rank=0),
parallel_config=SimpleNamespace(
enable_elastic_ep=False,
nnodes=1,
pipeline_parallel_size=1,
data_parallel_rank=0,
),
),
get_model=Mock(
return_value=SimpleNamespace(modules=Mock(return_value=modules))
@@ -133,6 +140,258 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn):
]
@pytest.mark.parametrize(
"has_v2,elastic_ep,node_count",
[(False, False, 1), (True, False, 1), (True, True, 1), (True, False, 2)],
ids=["v1", "v2", "elastic-v1", "multi-node-v1"],
)
@pytest.mark.parametrize(
"hisparse,world_size,rank,skip_ops,failure",
[
(False, 1, 0, [], None),
(False, 2, 0, [], None),
(False, 2, 1, [], None),
(True, 2, 0, ["fp4_gemm", "bmm_fp8"], None),
(False, 2, 1, ["fp4_gemm"], "dummy"),
(False, 2, 1, [], "replayssm"),
(True, 2, 0, [], "kimi"),
],
)
def test_flashinfer_autotune_lifecycle(
monkeypatch,
tmp_path,
has_v2,
elastic_ep,
node_count,
hisparse,
world_size,
rank,
skip_ops,
failure,
):
use_v2 = has_v2 and not elastic_ep and node_count == 1 and world_size == 1
runner = _make_runner([])
runner.vllm_config.parallel_config.enable_elastic_ep = elastic_ep
runner.vllm_config.attention_config.hisparse_config = object() if hisparse else None
model = runner.get_model()
cache_path = tmp_path / "autotune.json"
rank_cache_path = cache_path.with_name(f"autotune_dp0_rank{rank}.json")
rank_cache_path.write_text(json.dumps({"_metadata": {}, "op": 1}))
events = Mock()
events.tuner.load_configs.return_value = True
events.tuner._dirty = True
world = SimpleNamespace(
rank_in_group=rank,
world_size=world_size,
cpu_group=object(),
broadcast_object=events.broadcast,
barrier=events.barrier,
)
@contextmanager
def autotune_context(**kwargs):
assert torch.is_inference_mode_enabled()
events.enter()
try:
yield
finally:
events.exit()
events.autotune.side_effect = autotune_context
if failure is not None:
getattr(events, failure).side_effect = RuntimeError("warmup failed")
monkeypatch.setitem(
sys.modules,
"flashinfer.autotuner",
SimpleNamespace(
AutoTuner=SimpleNamespace(get=lambda: events.tuner),
set_autotune_process_group=events.set_group,
),
)
monkeypatch.setitem(
sys.modules,
"flashinfer",
SimpleNamespace(autotune_v2=events.autotune, autotune_v2_reload=events.reload),
)
monkeypatch.setattr(
"vllm.distributed.parallel_state.get_world_group", lambda: world
)
monkeypatch.setattr(
"vllm.distributed.parallel_state.get_node_count", lambda: node_count
)
monkeypatch.setattr(
"vllm.distributed.parallel_state.get_pp_group",
lambda: SimpleNamespace(world_size=1),
)
monkeypatch.setattr(fi_utils, "has_flashinfer_autotune_v2", lambda: has_v2)
monkeypatch.setattr(fi_utils, "autotune", events.autotune)
monkeypatch.setattr(warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS", skip_ops)
resolve_file = Mock(return_value=cache_path)
resolve_root = Mock(return_value=str(tmp_path))
monkeypatch.setattr(warmup, "resolve_flashinfer_autotune_file", resolve_file)
monkeypatch.setattr(warmup, "resolve_flashinfer_autotune_v2_root", resolve_root)
def gather(out, value, *, group):
out[:] = [value] * world_size
monkeypatch.setattr(torch.distributed, "all_gather_object", gather)
monkeypatch.setattr(warmup, "_run_flashinfer_autotune_dummy_runs", events.dummy)
monkeypatch.setattr(warmup, "replayssm_autotune_warmup", events.replayssm)
monkeypatch.setattr(warmup, "_autotune_kimi_k3_kda_qkvg", events.kimi)
monkeypatch.setattr(
warmup, "autotune_hisparse_flashinfer_attention", events.hisparse
)
if failure is None:
warmup.flashinfer_autotune(runner)
else:
with pytest.raises(RuntimeError, match="warmup failed"):
warmup.flashinfer_autotune(runner)
kwargs = {"skip_ops": set(skip_ops)} if skip_ops else {}
expected = []
if use_v2:
resolve_file.assert_not_called()
resolve_root.assert_called_once_with()
expected.append(call.autotune(mode="tune", cache_root=str(tmp_path), **kwargs))
else:
resolve_root.assert_not_called()
resolve_file.assert_called_once_with(runner)
expected.extend(
[
call.tuner.load_configs(str(rank_cache_path)),
call.autotune(tune_mode=True, **kwargs),
]
)
expected.extend(
[call.set_group(world.cpu_group if world_size > 1 else None), call.enter()]
)
if hisparse:
expected.append(call.hisparse(runner))
for name, arg, kwargs in [
("dummy", runner, {"skip_attn": hisparse}),
("replayssm", runner, {}),
("kimi", model, {}),
]:
expected.append(getattr(call, name)(arg, **kwargs))
if failure == name:
break
expected.append(call.exit())
expected.append(call.set_group(None))
if failure is None:
if world_size > 1:
expected.append(call.barrier())
if use_v2 and world_size > 1:
expected.extend([call.reload(), call.barrier()])
elif not use_v2:
expected.append(call.tuner.save_configs(str(rank_cache_path)))
assert events.mock_calls == expected
@pytest.mark.parametrize(
"node_count,elastic_ep,pp_size",
[(1, False, 1), (2, False, 1), (1, True, 1), (1, False, 2)],
)
@pytest.mark.parametrize("world_size,rank", [(1, 0), (2, 0), (2, 1)])
@pytest.mark.parametrize("model_runner_v2,max_reqs", [(False, 2), (True, 1), (True, 2)])
def test_sparse_autotune_respects_cache_selection(
monkeypatch,
tmp_path,
node_count,
elastic_ep,
pp_size,
world_size,
rank,
model_runner_v2,
max_reqs,
):
import vllm.model_executor.warmup.flashinfer_sparse_mla_warmup as sparse
runner = _make_runner([])
runner.max_num_reqs = max_reqs
runner.vllm_config.use_v2_model_runner = model_runner_v2
runner.vllm_config.kernel_config = SimpleNamespace(enable_flashinfer_autotune=True)
runner.vllm_config.parallel_config.enable_elastic_ep = elastic_ep
runner.vllm_config.parallel_config.pipeline_parallel_size = pp_size
worker = SimpleNamespace(
model_runner=runner,
vllm_config=runner.vllm_config,
execute_model=Mock(),
sample_tokens=Mock(),
)
world = SimpleNamespace(
rank_in_group=rank,
world_size=world_size,
broadcast_object=Mock(return_value=b"cached"),
barrier=Mock(),
)
managed = Mock(return_value=nullcontext())
legacy = Mock(return_value=nullcontext())
reload = Mock()
tuner = Mock()
monkeypatch.setitem(
sys.modules,
"flashinfer",
SimpleNamespace(autotune_v2=managed, autotune_v2_reload=reload),
)
monkeypatch.setitem(
sys.modules,
"flashinfer.autotuner",
SimpleNamespace(AutoTuner=SimpleNamespace(get=lambda: tuner)),
)
monkeypatch.setattr(
"vllm.distributed.parallel_state.get_world_group", lambda: world
)
monkeypatch.setattr(
"vllm.distributed.parallel_state.get_node_count", lambda: node_count
)
monkeypatch.setattr(fi_utils, "has_flashinfer_autotune_v2", lambda: True)
monkeypatch.setattr(sparse, "has_flashinfer", lambda: True)
monkeypatch.setattr(
sparse.current_platform, "is_device_capability_family", lambda _: True
)
monkeypatch.setattr(
sparse, "_flashinfer_sparse_mla_decode_label", lambda *_: "DSv3.2"
)
monkeypatch.setattr(sparse, "resolve_flashinfer_autotune_v2_root", lambda: tmp_path)
cache_path = tmp_path / "autotune.json"
cache_path.write_bytes(b"cached")
monkeypatch.setattr(
sparse, "resolve_flashinfer_autotune_file", lambda _: cache_path
)
monkeypatch.setattr(sparse, "flashinfer_autotune", legacy)
def mixed_warmup(*args, mixed_step_context=None, **kwargs):
context = (
mixed_step_context if mixed_step_context is not None else nullcontext()
)
with context:
assert torch.is_inference_mode_enabled()
return True
mixed = Mock(side_effect=mixed_warmup)
monkeypatch.setattr(sparse, "run_mixed_prefill_decode_warmup", mixed)
assert sparse._run_flashinfer_sparse_mla_decode_autotune(worker, 16, frozenset())
if node_count == 1 and not elastic_ep and pp_size == 1 and world_size == 1:
managed.assert_called_once_with(mode="tune", cache_root=tmp_path)
legacy.assert_not_called()
world.broadcast_object.assert_not_called()
tuner.load_configs.assert_not_called()
reload.assert_not_called()
world.barrier.assert_not_called()
else:
managed.assert_not_called()
reload.assert_not_called()
assert legacy.call_count == (1 if rank == 0 else 0)
world.broadcast_object.assert_called_once_with(
b"cached" if rank == 0 else None, src=0
)
tuner.load_configs.assert_called_once_with(str(cache_path))
assert mixed.call_count == (1 if model_runner_v2 and max_reqs >= 2 else 0)
assert runner._dummy_run.call_count == (0 if mixed.called else 1)
def test_flashinfer_autotune_buckets_cover_drafter_tokens():
"""The drafter's M can exceed the pass size; buckets must include it."""
runner = _make_runner([])
@@ -266,7 +525,9 @@ class _AutotuneRun:
for rank in range(self.pp * self.tp):
self.rank = rank
self.tuners[rank] = _AutotuneTuner(self)
flashinfer_autotune(_make_runner([]))
runner = _make_runner([])
runner.vllm_config.parallel_config.pipeline_parallel_size = self.pp
flashinfer_autotune(runner)
return self
def assert_collectives_match(self):
@@ -283,6 +544,8 @@ def autotune_run(monkeypatch, tmp_path):
def make_run(*, pp=2, tp=4):
run = _AutotuneRun(pp, tp, tmp_path)
monkeypatch.setattr(fi_utils, "has_flashinfer_autotune_v2", lambda: pp > 1)
monkeypatch.setattr(parallel_state, "get_node_count", lambda: 1)
autotuner = ModuleType("flashinfer.autotuner")
monkeypatch.setattr(
autotuner,
tests/model_executor/test_jit_warmup_triton_launcher.py
+17 -11
View File
@@ -168,23 +168,28 @@ def _warmup_store_kernel(out, VALUE: tl.constexpr):
not current_platform.is_cuda() or not hasattr(triton, "AsyncCompileMode"),
reason="Requires CUDA Triton async compilation",
)
def test_registered_warmup_parallelizes_compile_only_and_runtime_stays_serial(
def test_registered_warmup_overlaps_kernels_and_runtime_stays_serial(
monkeypatch,
) -> None:
from triton.runtime._async_compile import active_mode
monkeypatch.setattr(envs, "VLLM_TRITON_JIT_WARMUP_NUM_THREADS", 2)
def warmup_inputs():
return dict(out=TritonWarmupTensor(torch.int32), value=WarmupChoices(1, 2))
def make_store(value):
def warmup_inputs():
return dict(out=TritonWarmupTensor(torch.int32), value=value)
@triton_kernel_dispatcher_with_warmup(
kernel=_warmup_store_kernel, warmup_inputs=warmup_inputs
)
def store(out, value):
return (1,), dict(VALUE=value)
@triton_kernel_dispatcher_with_warmup(
kernel=_warmup_store_kernel, warmup_inputs=warmup_inputs
)
def store(out, value):
return (1,), dict(VALUE=value)
store.get_warmup_keys() # Initialize Triton's binder before wrapping its compiler.
# Initialize Triton's binder before wrapping its compiler.
store.get_warmup_keys()
return store
stores = [make_store(value) for value in (1, 2)]
caller = threading.get_ident()
barrier = threading.Barrier(2, timeout=10)
compiled_on = []
@@ -200,14 +205,15 @@ def test_registered_warmup_parallelizes_compile_only_and_runtime_stays_serial(
monkeypatch.setattr(_warmup_store_kernel, "compile", record_compile)
registry = JitWarmupRegistry(None)
with registry.activate():
store.register_warmup()
for store in stores:
store.register_warmup()
registry.warmup()
assert len(compiled_on) == 2
assert all(thread != caller for thread in compiled_on)
assert active_mode.get() is None
out = torch.empty(1, dtype=torch.int32, device=current_platform.device_type)
for value in (1, 2):
for store, value in zip(stores, (1, 2)):
store(out, value)
assert out.item() == value
assert len(compiled_on) == 2
tests/model_executor/test_triton_dispatcher.py
+41
View File
@@ -220,3 +220,44 @@ def test_real_triton_kernel_override_end_to_end() -> None:
"out_ptr",
"BLOCK",
)
def test_cpu_slot_mapping_override_forwards_kernel_arguments(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm.utils import cpu_triton_utils
from vllm.v1.worker.block_table import ComputeSlotMappingKernel
calls: list[tuple[Any, ...]] = []
fake_torch = types.SimpleNamespace(
ops=types.SimpleNamespace(
_C=types.SimpleNamespace(
compute_slot_mapping_kernel_impl=lambda *args: calls.append(args)
)
)
)
monkeypatch.setattr(cpu_triton_utils, "torch", fake_torch)
override = KernelOverride(
ComputeSlotMappingKernel.kernel,
cpu_triton_utils._compute_slot_mapping_kernel_impl,
)
override[(2,)](
num_tokens=2,
max_num_tokens=4,
query_start_loc_ptr="query_start_loc",
positions_ptr="positions",
block_table_ptr="block_table",
block_table_stride=128,
slot_mapping_ptr="slot_mapping",
KV_CACHE_BLOCK_SIZE=16,
TOTAL_CP_WORLD_SIZE=1,
TOTAL_CP_RANK=0,
CP_KV_CACHE_INTERLEAVE_SIZE=1,
PAD_ID=-1,
BLOCK_SIZE=1024,
)
assert calls == [
("query_start_loc", "positions", "block_table", "slot_mapping", 16)
]
tests/models/kimi_k3/test_amd_kda_adaptive_verification.py
+93
View File
@@ -0,0 +1,93 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""ROCm Kimi-K3 KDA metadata under adaptive verification.
Adaptive verification splits the draft budget evenly on the host, then trims
each verify request on device. The KDA spec path must therefore plan from the
device offsets, never from host lengths that merely look uniform.
"""
import torch
from tests.v1.attention.utils import (
BatchSpec,
create_common_attn_metadata,
create_vllm_config,
)
from vllm.config import SpeculativeConfig
from vllm.config.compilation import CUDAGraphMode
from vllm.models.kimi_k3.amd.kda_metadata import (
KimiK3ROCmKDABackend,
KimiK3ROCmKDAMetadataBuilder,
)
from vllm.v1.kv_cache_interface import MambaSpec
BLOCK_SIZE = 16
DEVICE = torch.device("cpu")
NUM_SPEC = 3
def _config(adaptive: bool):
vllm_config = create_vllm_config(
model_name="Qwen/Qwen3.5-0.8B",
block_size=BLOCK_SIZE,
)
vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
vllm_config.speculative_config = SpeculativeConfig(
method="ngram",
num_speculative_tokens=NUM_SPEC,
enable_adaptive_verification=adaptive,
)
vllm_config.cache_config.mamba_cache_mode = "none"
return vllm_config
def _spec() -> MambaSpec:
return MambaSpec(
block_size=BLOCK_SIZE,
shapes=((16, 64),),
dtypes=(torch.float16,),
num_speculative_blocks=NUM_SPEC,
)
def _builder(adaptive: bool) -> KimiK3ROCmKDAMetadataBuilder:
return KimiK3ROCmKDAMetadataBuilder(
kv_cache_spec=_spec(),
layer_names=["layer.0"],
vllm_config=_config(adaptive),
device=DEVICE,
)
def _trimmed_verify_batch(builder: KimiK3ROCmKDAMetadataBuilder):
"""Host sees [3, 3]; the device was trimmed to [4, 2] (same total)."""
batch = BatchSpec(seq_lens=[64, 64], query_lens=[4, 2])
common = create_common_attn_metadata(
batch, BLOCK_SIZE, DEVICE, arange_block_indices=True
)
common = common.replace(
query_start_loc_cpu=torch.tensor([0, 3, 6], dtype=torch.int32)
)
accepted = torch.ones(2, dtype=torch.int32, device=DEVICE)
num_decode_draft_tokens_cpu = torch.tensor([NUM_SPEC, NUM_SPEC], dtype=torch.int32)
return builder.build(0, common, accepted, num_decode_draft_tokens_cpu)
def test_adaptive_verification_plans_from_device_offsets():
"""Equal host lengths would select the fixed-length recurrent kernel,
which places sequence i at i * L and ignores the device cu_seqlens."""
assert KimiK3ROCmKDABackend.supports_device_cpu_query_lens_mismatch()
bound = KimiK3ROCmKDAMetadataBuilder.get_varlen_cudagraph_max_query_len
assert bound(_config(adaptive=True), _spec()) == NUM_SPEC + 1
assert bound(_config(adaptive=False), _spec()) is None
meta = _trimmed_verify_batch(_builder(adaptive=True))
assert meta.num_spec_decodes == 2
assert meta.uniform_spec_sequence_length is None
assert torch.equal(
meta.spec_query_start_loc.cpu(), torch.tensor([0, 4, 6], dtype=torch.int32)
)
# Without adaptive verification the host lengths are exact and kept.
meta = _trimmed_verify_batch(_builder(adaptive=False))
assert meta.uniform_spec_sequence_length == 3
tests/models/kimi_k3/test_amd_kda_checkpoint.py
+3 -2
View File
@@ -482,7 +482,7 @@ def _spec_for(monkeypatch, **attrs) -> MambaSpec:
layer = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention)
for name, value in dict(
use_prefill_checkpoint=True,
use_fused_chunk=True,
kda_prefill_backend="fused",
num_spec=0,
use_safe_gate=True,
).items():
@@ -506,7 +506,8 @@ def test_opt_in_publishes_chunk_walk_alignment(monkeypatch) -> None:
"attr,value",
[
("use_prefill_checkpoint", False),
("use_fused_chunk", False),
("kda_prefill_backend", "triton"),
("kda_prefill_backend", "gluon"),
("use_safe_gate", False),
],
)
tests/models/kimi_k3/test_amd_kda_chunk.py
+23 -2
View File
@@ -77,7 +77,7 @@ def _run(inp: dict, use_fused: bool, out: torch.Tensor | None = None, **kw):
output_final_state=True,
use_qk_l2norm_in_kernel=True,
cu_seqlens=inp["cu"],
use_fused_chunk=use_fused,
backend=kw.pop("backend", "fused" if use_fused else "triton"),
out=out,
**kw,
)
@@ -112,6 +112,27 @@ def test_fused_chunk_matches_triton(seqlens: list[int]) -> None:
torch.testing.assert_close(state_fused, state_ref, rtol=1e-3, atol=1e-3)
@pytest.mark.parametrize(
"seqlens",
[[1], [64], [1024], [4096], [513, 64, 1, 1200], [512] * 6],
)
def test_gluon_chunk_matches_triton(seqlens: list[int]) -> None:
"""AITER's Gluon chunk_kda must match the Triton chunk path too."""
from vllm._aiter_ops import rocm_aiter_ops
if not rocm_aiter_ops.is_kda_gluon_prefill_available():
pytest.skip("AITER Gluon chunk_kda is not available here")
inp = _inputs(seqlens, seed=len(seqlens))
o_ref, state_ref = _run(inp, use_fused=False)
o_gluon, state_gluon = _run(inp, use_fused=False, backend="gluon")
torch.testing.assert_close(o_gluon.float(), o_ref.float(), rtol=2e-2, atol=2e-2)
torch.testing.assert_close(state_gluon, state_ref, rtol=5e-3, atol=5e-3)
state_rel = (state_gluon - state_ref).norm() / state_ref.norm()
assert state_rel < 5e-3, f"final state relative L2 error {state_rel:.2e}"
@pytest.mark.parametrize("use_fused", [False, True])
def test_out_argument_receives_the_result(use_fused: bool) -> None:
"""``out`` is a promise about where the result lands.
@@ -648,7 +669,7 @@ def _run_paged(inp: dict, cache, idx, warm, use_fused: bool, **kw):
has_initial_state=warm,
use_qk_l2norm_in_kernel=True,
cu_seqlens=inp["cu"],
use_fused_chunk=use_fused,
backend="fused" if use_fused else "triton",
**kw,
)
tests/models/kimi_k3/test_amd_kda_fused_conv1d.py
+163
View File
@@ -0,0 +1,163 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_fn
from vllm.platforms import current_platform
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
pytestmark = pytest.mark.skipif(
not current_platform.is_rocm(), reason="AITER fused conv1d is ROCm-only"
)
WIDTH = 4 # Kimi-K3 short_conv_kernel_size
STATE_LEN = WIDTH - 1
HEAD_DIM = 128
LOCAL_HEADS = 4 # per-TP-rank heads; projection = 512, packed dim = 1536
def _fused_fn():
gdn = pytest.importorskip("aiter.ops.triton.gated_delta_net")
return gdn.causal_conv1d_split_qkv_triton_tile_fn
def _make_inputs(seqlens, has_initial_state, seed=0, device="cuda"):
torch.manual_seed(seed)
proj = HEAD_DIM * LOCAL_HEADS
dim = 3 * proj
total = sum(seqlens)
n_seqs = len(seqlens)
# Token-major activations, as produced by the in_proj GEMM.
mixed_qkv = torch.randn(total, dim, dtype=torch.bfloat16, device=device)
weight = torch.randn(dim, WIDTH, dtype=torch.bfloat16, device=device) * 0.3
# More cache lines than sequences, assigned in a shuffled order. Line 0 is
# vLLM's null block (NULL_BLOCK_ID), which causal_conv1d_fn skips, so real
# sequences never map to it.
n_lines = n_seqs + 2
conv_state = torch.randn(
n_lines, dim, STATE_LEN, dtype=torch.bfloat16, device=device
)
cache_indices = (torch.randperm(n_lines - 1, device=device)[:n_seqs] + 1).to(
torch.int32
)
query_start_loc = torch.tensor(
[0] + list(torch.tensor(seqlens).cumsum(0).tolist()),
dtype=torch.int32,
device=device,
)
has_initial = torch.full(
(n_seqs,), has_initial_state, dtype=torch.bool, device=device
)
return mixed_qkv, weight, conv_state, cache_indices, query_start_loc, has_initial
def _reference(
mixed_qkv, weight, conv_state, cache_indices, query_start_loc, has_initial
):
"""fp32 per-sequence depthwise causal conv + SiLU.
Returns (out[T, dim], new_conv_state).
"""
x = mixed_qkv.float().t() # (dim, T)
w = weight.float()
dim = x.shape[0]
out = torch.empty_like(x)
new_state = conv_state.float().clone()
qsl = query_start_loc.tolist()
for s in range(len(qsl) - 1):
start, end = qsl[s], qsl[s + 1]
ci = int(cache_indices[s])
xs = x[:, start:end]
if bool(has_initial[s]):
xin = torch.cat([conv_state[ci].float(), xs], dim=1)
else:
xin = F.pad(xs, (STATE_LEN, 0))
y = F.conv1d(xin.unsqueeze(0), w.unsqueeze(1), None, groups=dim)[0]
out[:, start:end] = F.silu(y)
new_state[ci] = xin[:, -STATE_LEN:]
return out.t(), new_state
def _three_launch(
mixed_qkv, weight, conv_state, cache_indices, query_start_loc, has_initial
):
"""The pre-fusion KDA prefill path: one causal_conv1d_fn per projection."""
proj = mixed_qkv.shape[1] // 3
outs = []
for x, state, w in zip(
mixed_qkv.split(proj, dim=-1),
conv_state.split(proj, dim=-2),
weight.split(proj, dim=0),
):
outs.append(
causal_conv1d_fn(
x.transpose(0, 1),
w,
None,
activation="silu",
conv_states=state,
has_initial_state=has_initial,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
).transpose(0, 1)
)
return outs
@pytest.mark.parametrize(
"seqlens", [[1024], [8, 1, 300], [64, 64, 64], [513, 64, 1, 1200]]
)
@pytest.mark.parametrize("has_initial_state", [False, True])
@pytest.mark.parametrize("block_m", [8, 128])
def test_fused_split_qkv_conv1d_matches_reference(seqlens, has_initial_state, block_m):
fused_fn = _fused_fn()
mixed_qkv, weight, conv_state, cache_indices, qsl, has_initial = _make_inputs(
seqlens, has_initial_state
)
proj = mixed_qkv.shape[1] // 3
ref_out, ref_state = _reference(
mixed_qkv, weight, conv_state, cache_indices, qsl, has_initial
)
state_3 = conv_state.clone()
q3, k3, v3 = _three_launch(
mixed_qkv, weight, state_3, cache_indices, qsl, has_initial
)
state_f = conv_state.clone()
# Exactly the call made by KimiK3DeltaAttention._forward.
qf, kf, vf = fused_fn(
mixed_qkv.transpose(0, 1),
weight,
None,
state_f,
qsl,
proj,
proj,
cache_indices=cache_indices,
has_initial_state=has_initial,
activation="silu",
block_m=block_m,
metadata=None,
preserve_input_dtype=True,
)
fused = torch.cat([qf, kf, vf], dim=-1)
assert fused.shape == mixed_qkv.shape and fused.dtype == mixed_qkv.dtype
assert all(t.is_contiguous() for t in (qf, kf, vf))
tol = dict(atol=2e-2, rtol=2e-2)
torch.testing.assert_close(fused.float(), ref_out, **tol)
torch.testing.assert_close(fused, torch.cat([q3, k3, v3], dim=-1), **tol)
# conv_state writeback: assigned lines updated, the others untouched.
torch.testing.assert_close(state_f.float(), ref_state, **tol)
torch.testing.assert_close(state_f, state_3, **tol)
untouched = torch.ones(conv_state.shape[0], dtype=torch.bool, device="cuda")
untouched[cache_indices.long()] = False
assert untouched[NULL_BLOCK_ID]
assert torch.equal(state_f[untouched], conv_state[untouched])
tests/multimodal/media/test_image.py
+18
View File
@@ -85,6 +85,24 @@ def test_image_media_io_rgba_custom_background(tmp_path):
assert green_numpy[0][0][2] == 0 # B
@pytest.mark.parametrize("background_color", [(255, 255, 255), (12, 34, 56)])
def test_image_media_io_rgb_png_transparency(tmp_path, background_color):
image = Image.new("RGB", (2, 1), (255, 0, 0))
image.putpixel((1, 0), (0, 0, 255))
path = tmp_path / "rgb_transparency.png"
image.save(path, transparency=(255, 0, 0))
with Image.open(path) as source:
reference = Image.alpha_composite(
Image.new("RGBA", source.size, (*background_color, 255)),
source.convert("RGBA"),
).convert("RGB")
image_io = ImageMediaIO(rgba_background_color=background_color)
converted = image_io.load_file(path)
np.testing.assert_array_equal(np.asarray(converted), np.asarray(reference))
def test_image_media_io_no_mode_conversion(tmp_path):
"""image_mode=None skips conversion and preserves the original mode."""
# RGBA image: opaque black pixel on a fully transparent background
tests/multimodal/test_hasher.py
+10 -2
View File
@@ -237,8 +237,16 @@ def _rgba_png_bytes() -> bytes:
return buf.getvalue()
def test_hash_collision_media_io_config():
data = _rgba_png_bytes()
def _rgb_trns_png_bytes() -> bytes:
image = Image.new("RGB", (8, 8), (255, 0, 0))
buf = BytesIO()
image.save(buf, format="PNG", transparency=(255, 0, 0))
return buf.getvalue()
@pytest.mark.parametrize("png_bytes", [_rgba_png_bytes, _rgb_trns_png_bytes])
def test_hash_collision_media_io_config(png_bytes):
data = png_bytes()
white = ImageMediaIO(rgba_background_color=(255, 255, 255)).load_bytes(data)
black = ImageMediaIO(rgba_background_color=(0, 0, 0)).load_bytes(data)
white2 = ImageMediaIO(rgba_background_color=(255, 255, 255)).load_bytes(data)
tests/multimodal/test_image.py
+65
View File
@@ -5,6 +5,7 @@ from io import BytesIO
from pathlib import Path
import numpy as np
import pybase64
import pytest
from PIL import Image, ImageChops
@@ -13,6 +14,8 @@ from vllm.multimodal.image import (
convert_image_mode,
normalize_image,
)
from vllm.multimodal.media import MediaConnector
from vllm.multimodal.parse import MultiModalDataParser
pytestmark = pytest.mark.cpu_test
@@ -69,6 +72,68 @@ def test_palette_with_trns_to_rgb():
assert (r, g, b) == (255, 0, 0)
def _reopen(image: Image.Image, image_format: str, **save_kwargs) -> Image.Image:
buffer = BytesIO()
image.save(buffer, image_format, **save_kwargs)
return Image.open(buffer)
def _rgb_png() -> Image.Image:
image = Image.new("RGB", (2, 1), (255, 0, 0))
image.putpixel((1, 0), (0, 0, 255))
return _reopen(image, "PNG", transparency=(255, 0, 0))
def _gif() -> Image.Image:
# The transparent index 0 and the opaque index 1 are both black.
image = Image.new("P", (2, 1))
image.putpalette([0] * 768)
image.putpixel((1, 0), 1)
return _reopen(image, "GIF", transparency=0)
def _rgba_png() -> Image.Image:
image = Image.new("RGBA", (2, 1), (255, 0, 0, 0))
image.putpixel((1, 0), (0, 0, 255, 255))
return _reopen(image, "PNG")
@pytest.mark.parametrize(
"make_image",
[
_rgb_png,
lambda: _rgb_png().convert("RGB"),
_gif,
lambda: _gif().convert("RGB"),
_rgba_png,
],
ids=["rgb_png", "rgb_png_to_rgb", "gif", "gif_to_rgb", "rgba_png"],
)
def test_transparency_offline_matches_online(make_image):
"""A PIL image passed offline and the same image sent online as a PNG
end up with the same pixels, matching Pillow's own RGBA conversion.
After ``convert("RGB")``, Pillow keeps the transparent color in ``info``
and saves it as a tRNS chunk, so every pixel of that color is
transparent in both pipelines, including opaque GIF pixels that share it.
"""
image = make_image()
buffer = BytesIO()
image.save(buffer, "PNG")
url = f"data:image/png;base64,{pybase64.b64encode(buffer.getvalue()).decode()}"
parser = MultiModalDataParser()
offline = parser.parse_mm_data({"image": image})["image"].get(0)
online = parser.parse_mm_data({"image": MediaConnector().fetch_image(url)})
online = online["image"].get(0)
np.testing.assert_array_equal(np.asarray(offline), np.asarray(online))
expected = Image.alpha_composite(
Image.new("RGBA", image.size, (255, 255, 255, 255)), image.convert("RGBA")
).convert("RGB")
np.testing.assert_array_equal(np.asarray(online), np.asarray(expected))
def test_l_mode_no_trns_to_rgb():
"""L-mode without transparency should convert directly."""
img = Image.new("L", (4, 4), 128)
tests/parser/engine/test_minimax_m2.py
+121
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import dataclasses
import json
import pytest
@@ -15,11 +16,14 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionToolsParam,
FunctionDefinition,
)
from vllm.parser.engine.parser_engine import ParserEngine
from vllm.parser.minimax_m2 import (
THINK_END,
TOOL_CALL_END,
TOOL_CALL_START,
MinimaxM2Parser,
_minimax_m2_arg_value_extender,
minimax_m2_config,
)
@@ -238,6 +242,123 @@ class TestStreaming:
assert collect_tool_arguments(results) == ""
WRITE_CALL_HEAD = (
'<minimax:tool_call>\n<invoke name="write_file">\n'
'<parameter name="path">/repo/app.py</parameter>\n'
'<parameter name="mode">644</parameter>\n'
'<parameter name="content">'
)
WRITE_CALL_TAIL = "</parameter>\n</invoke>\n</minimax:tool_call>"
# Escapes, non-ASCII, blank lines and tag-like text that must stay inside
# the value, including prefixes of the closing and next-parameter tags.
CONTENT_LINE = 'if a < b > c: print("x\\y") # </p> </param <parameter\té\n\n'
class TestArgValueExtender:
"""The open parameter value streams without re-running the converter."""
@pytest.fixture
def write_tool(self):
return ChatCompletionToolsParam(
function=FunctionDefinition(
name="write_file",
parameters={
"type": "object",
"properties": {
"path": {"type": "string"},
"mode": {"type": "integer"},
"content": {"type": "string"},
},
},
),
)
@pytest.mark.parametrize("use_schema", [False, True], ids=["no_schema", "schema"])
def test_long_call_streams_non_streaming_arguments(
self, mock_tokenizer, mock_request, write_tool, use_schema
):
mock_request.tools = [write_tool] if use_schema else []
text = WRITE_CALL_HEAD + CONTENT_LINE * 200 + WRITE_CALL_TAIL
results = simulate_tool_streaming(
MinimaxM2Parser(mock_tokenizer),
mock_request,
[text[i : i + 7] for i in range(0, len(text), 7)],
)
result = MinimaxM2Parser(mock_tokenizer).extract_tool_calls(text, mock_request)
arguments = collect_tool_arguments(results)
assert arguments == result.tool_calls[0].function.arguments
assert json.loads(arguments)["content"] == CONTENT_LINE * 200
assert ('"mode": 644,' if use_schema else '"mode": "644",') in arguments
def test_conversions_do_not_grow_with_argument_length(
self, mock_tokenizer, mock_request, write_tool
):
mock_request.tools = [write_tool]
base = minimax_m2_config()
converter = base.arg_converter
assert converter is not None
def streaming_conversions(lines: int) -> int:
calls = 0
def counting_converter(raw_args: str, partial: bool) -> str:
nonlocal calls
calls += partial
return converter(raw_args, partial)
parser = ParserEngine(
mock_tokenizer,
parser_engine_config=dataclasses.replace(
base, arg_converter=counting_converter
),
)
text = WRITE_CALL_HEAD + CONTENT_LINE * lines + WRITE_CALL_TAIL
simulate_tool_streaming(
parser,
mock_request,
[text[i : i + 4] for i in range(0, len(text), 4)],
)
return calls
assert streaming_conversions(lines=20) == streaming_conversions(lines=200)
def test_extender_stops_at_split_end_tag(self):
extend = _minimax_m2_arg_value_extender('<parameter name="k">')
assert extend is not None
deltas = ['a"\n', "</param", "eter>"]
assert [extend(d) for d in deltas] == ['a\\"\\n', "</param", None]
def test_extender_stops_at_split_next_parameter(self):
extend = _minimax_m2_arg_value_extender("<parameter name='k'>x")
assert extend is not None
assert [extend(d) for d in ["< parameter", " name ="]] == ["< parameter", None]
@pytest.mark.parametrize(
"raw_args",
[
'<parameter name="k">x</parameter>',
'<parameter name="k',
(
'<parameter name="k">1</parameter><parameter name="j">2</parameter>'
'<parameter name="k">3'
),
'<parameter name="a<parameter name=k>x',
'<parameter name="k>x',
],
ids=[
"closed_value",
"open_name",
"repeated_key_not_last",
"earlier_unmatched_opener",
"bare_name_opens_quote",
],
)
def test_no_extender_unless_open_value_is_stable(self, raw_args):
assert _minimax_m2_arg_value_extender(raw_args) is None
class TestReasoning:
def test_extract_reasoning_without_start_token(self, parser, mock_request):
reasoning, content = parser.extract_reasoning(
tests/parser/engine/test_qwen3_reasoning.py
+34
View File
@@ -380,6 +380,40 @@ class TestDelegatingPromptDetection:
assert delta.reasoning is None
assert delta.content == "answer"
def test_prior_turn_think_end_keeps_non_streaming_reasoning(
self, mock_tokenizer, mock_request
):
parser = _Qwen3DelegatingParser(mock_tokenizer)
parser.set_prompt_token_ids([_THINK_END_ID, _IM_END_ID, _IM_START_ID, _TEXT_ID])
output_ids = [ord("a"), _THINK_END_ID, ord("b")]
reasoning, content, _ = parser.parse(
mock_tokenizer.decode(output_ids),
mock_request,
model_output_token_ids=output_ids,
)
assert reasoning == "a"
assert content == "b"
assert parser.count_reasoning_tokens(output_ids) > 0
def test_think_end_in_current_turn_parses_output_as_content(
self, mock_tokenizer, mock_request
):
parser = _Qwen3DelegatingParser(mock_tokenizer)
parser.set_prompt_token_ids([_IM_START_ID, _THINK_START_ID, _THINK_END_ID])
output_ids = [ord("a"), ord("b")]
reasoning, content, _ = parser.parse(
mock_tokenizer.decode(output_ids),
mock_request,
model_output_token_ids=output_ids,
)
assert reasoning is None
assert content == "ab"
assert parser.count_reasoning_tokens(output_ids) == 0
class TestStreaming:
def test_basic_streaming(self, parser):
tests/parser/test_beam_parser.py
+747
View File
@@ -0,0 +1,747 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from copy import deepcopy
from types import SimpleNamespace
from typing import Any
import pytest
from openai.types.responses import FunctionTool, ToolChoiceFunction
from xgrammar import Grammar
from xgrammar.testing import _is_grammar_accept_string
from vllm.config.reasoning import ReasoningConfig
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
ChatCompletionToolsParam,
)
from vllm.entrypoints.openai.responses.protocol import (
ResponsesRequest,
ResponsesResponse,
)
from vllm.exceptions import VLLMValidationError
from vllm.parser.beam import BeamParser, BeamReasoningParser, BeamToolParser
from vllm.parser.parser_manager import ParserManager
from vllm.parser.response_template import ResponseTemplateParser
from vllm.reasoning import ReasoningParserManager
from vllm.tool_parsers import ToolParserManager
# The `response_template` of the Beam checkpoints.
BEAM_RESPONSE_TEMPLATE: dict[str, Any] = {
"defaults": {"role": "assistant"},
"start_anchor_pattern": r"<\|start\|>assistant",
"fields": {
"thinking": {
"open": "<|channel|>analysis<|message|>",
"close": "<|end|><|start|>assistant",
"content": "text",
},
"content": {
"open": "<|channel|>final<|message|>",
"close": ["<|return|>", "<|end|>", "<|call|>"],
"content": "text",
},
"tool_calls": {
"open_pattern": (
r" to=(?P<name>[^\s<]+)"
r"<\|channel\|>analysis <\|constrain\|>xml<\|message\|>"
),
"close": ["<|end|><|start|>assistant", "<|call|>"],
"repeats": True,
"content": "xml-inline",
"content_args": {
"tag_pattern": (
r'<rfl:parameter name="(?P<key>[^"]+)">'
r'\n(?P<value>.*?)\n(?:</rfl:parameter>|(?=<rfl:parameter name="))'
),
"value_parser": {"name": "text", "args": {"strip": False}},
},
"transform": {
"type": "function",
"function": {"name": "{name}", "arguments": "{content}"},
},
},
},
}
VOCAB = {
"<|start|>": 200,
"<|channel|>": 201,
"<|message|>": 202,
"<|end|>": 203,
"<|call|>": 204,
"<|return|>": 205,
"<|constrain|>": 206,
}
TEXT_TOKENS = {100: "assistant", 101: "analysis", 102: "final", 103: "plan"}
PROMPT_TOKEN_IDS = [200, 100]
THOUGHT = "<|channel|>analysis<|message|>plan<|end|><|start|>assistant"
FINAL = "<|channel|>final<|message|>"
NEXT_MESSAGE = "<|end|><|start|>assistant"
def _tool(name, properties, *, strict=None):
return ChatCompletionToolsParam(
type="function",
function={
"name": name,
"parameters": {"type": "object", "properties": properties},
"strict": strict,
},
)
TOOLS = [
_tool("get_weather", {"city": {"type": "string"}}),
_tool("set_count", {"count": {"type": "integer"}}),
]
STRICT_TOOLS = [
_tool("get_weather", {"city": {"type": "string"}}),
_tool("set_count", {"count": {"type": "integer"}}, strict=True),
]
class BeamTokenizer:
def __init__(self, response_template=BEAM_RESPONSE_TEMPLATE):
self.response_template = response_template
self.init_kwargs = {}
self.name_or_path = ""
self.all_special_tokens = list(VOCAB)
self.all_special_ids = list(VOCAB.values())
def decode(self, token_ids, skip_special_tokens=False, **_kwargs):
inverse_vocab = {token_id: token for token, token_id in VOCAB.items()}
return "".join(
TEXT_TOKENS.get(token_id, inverse_vocab.get(token_id, ""))
for token_id in token_ids
)
def get_vocab(self):
return VOCAB
def encode(self, text, add_special_tokens=False):
del add_special_tokens
return list(text.encode())
def request(**overrides):
values = dict(
tool_choice="auto",
include_reasoning=True,
tools=TOOLS,
parallel_tool_calls=True,
)
values.update(overrides)
return SimpleNamespace(**values)
def parameter(name, value):
return f'<rfl:parameter name="{name}">\n{value}\n</rfl:parameter>'
def call(name, body, *, terminator="<|call|>"):
return (
f" to={name}<|channel|>analysis <|constrain|>xml<|message|>{body}{terminator}"
)
def make_parser(prompt_token_ids=PROMPT_TOKEN_IDS, tools=TOOLS, **kwargs):
parser_cls = ParserManager.get_parser(
tool_parser_name="beam",
reasoning_parser_name="beam",
enable_auto_tools=True,
tokenizer=BeamTokenizer(),
)
assert parser_cls is not None
parser = parser_cls(BeamTokenizer(), tools, **kwargs)
parser.set_prompt_token_ids(prompt_token_ids)
return parser
def parse(text, prompt_token_ids=PROMPT_TOKEN_IDS):
reasoning, content, calls = make_parser(prompt_token_ids).parse(
text, request(), enable_auto_tools=True
)
return reasoning, content, normalize_calls(calls)
def normalize_calls(calls):
return [(call.name, json.loads(call.arguments)) for call in calls or []]
def test_beam_parsers_are_registered():
assert ToolParserManager.get_tool_parser("beam") is BeamToolParser
assert ReasoningParserManager.get_reasoning_parser("beam") is BeamReasoningParser
def test_parser_manager_selects_beam_parser():
parser_cls = ParserManager.get_parser(
tool_parser_name="beam",
reasoning_parser_name="beam",
enable_auto_tools=True,
tokenizer=BeamTokenizer(),
)
assert parser_cls is not None
assert issubclass(parser_cls, BeamParser)
assert parser_cls.reasoning_parser_cls is BeamReasoningParser
assert parser_cls.tool_parser_cls is BeamToolParser
assert parser_cls._enable_auto_tools
@pytest.mark.parametrize("other", ["hermes", "hf"])
def test_beam_parser_rejects_other_parsers(other):
with pytest.raises(TypeError, match="cannot be combined"):
ParserManager.get_parser(
tool_parser_name=other,
reasoning_parser_name="beam",
enable_auto_tools=True,
tokenizer=BeamTokenizer(),
)
def test_beam_parser_requires_checkpoint_metadata_at_startup():
with pytest.raises(
TypeError, match="--reasoning-parser beam requires `response_template`"
):
ParserManager.get_parser(
reasoning_parser_name="beam",
tokenizer=BeamTokenizer(response_template=None),
)
template = deepcopy(BEAM_RESPONSE_TEMPLATE)
del template["fields"]["tool_calls"]
with pytest.raises(TypeError, match="--tool-call-parser beam requires"):
ParserManager.get_parser(
tool_parser_name="beam",
enable_auto_tools=True,
tokenizer=BeamTokenizer(template),
)
def test_reasoning_and_final_answer():
assert parse(THOUGHT + FINAL + "Sunny.<|return|>") == ("plan", "Sunny.", [])
def test_prefilled_reasoning_continues_from_prompt():
prompt = PROMPT_TOKEN_IDS + [201, 101, 202, 103]
text = " more.<|end|><|start|>assistant" + FINAL + "Done.<|return|>"
assert parse(text, prompt) == (" more.", "Done.", [])
def test_parallel_calls_are_coerced_to_their_schemas():
text = (
THOUGHT
+ call(
"get_weather",
parameter("city", "Paris"),
terminator="<|end|><|start|>assistant",
)
+ call("set_count", parameter("count", "7"))
)
assert parse(text) == (
"plan",
None,
[("get_weather", {"city": "Paris"}), ("set_count", {"count": 7})],
)
def test_call_stop_token_removed_from_the_text_still_ends_the_call():
text = THOUGHT + call("get_weather", parameter("city", "Paris"), terminator="")
assert parse(text) == ("plan", None, [("get_weather", {"city": "Paris"})])
@pytest.mark.parametrize("end", ["<|end|>", "<|end|><|start|>"])
def test_output_cut_before_the_next_message_start_ends_its_region(end):
"""Output that stops after `<|end|>`, such as at `max_tokens`, has ended
its reasoning or call and lacks only the next message's start."""
weather = call("get_weather", parameter("city", "Paris"), terminator=end)
assert parse("<|channel|>analysis<|message|>plan" + end) == ("plan", None, [])
assert parse(THOUGHT + weather) == (
"plan",
None,
[("get_weather", {"city": "Paris"})],
)
THOUGHT_IDS = [201, 101, 202, 103, 203, 200, 100]
FINAL_IDS = [201, 102, 202]
@pytest.mark.parametrize(
("prompt_token_ids", "token_ids", "count"),
[
(PROMPT_TOKEN_IDS, [*THOUGHT_IDS, *FINAL_IDS, 205], 1),
(PROMPT_TOKEN_IDS + [201, 101, 202], [103, 103, 203, 200, 100, *FINAL_IDS], 2),
(PROMPT_TOKEN_IDS, [201, 101, 202, 103, 103, 103], 3),
],
ids=["reasoning_then_answer", "prefilled_reasoning", "cut_off"],
)
def test_reasoning_tokens_are_counted(prompt_token_ids, token_ids, count):
parser = make_parser(prompt_token_ids)
assert parser.count_reasoning_tokens(token_ids) == count
def test_reasoning_tokens_are_counted_as_the_generation_grows():
token_ids = [*THOUGHT_IDS, *FINAL_IDS, 205]
parser = make_parser()
assert [
parser.count_reasoning_tokens(token_ids[:end])
for end in range(1, len(token_ids) + 1)
] == [
make_parser().count_reasoning_tokens(token_ids[:end])
for end in range(1, len(token_ids) + 1)
]
assert parser.count_reasoning_tokens(FINAL_IDS) == 0
def test_call_stop_token_ends_the_answer():
assert parse(FINAL + "done<|call|>") == (None, "done", [])
@pytest.mark.parametrize("chunk_size", [1, 7, 10_000])
def test_streaming_matches_non_streaming(chunk_size):
text = (
THOUGHT
+ call(
"get_weather",
parameter("city", "Paris"),
terminator="<|end|><|start|>assistant",
)
+ call("set_count", parameter("count", "7"))
)
parser = make_parser()
reasoning, calls = [], []
for start in range(0, len(text), chunk_size):
end = min(start + chunk_size, len(text))
delta = parser.parse_delta(
text[start:end], [], request(), finished=end == len(text)
)
if delta is not None:
reasoning.append(delta.reasoning or "")
calls.extend(delta.tool_calls)
assert "".join(reasoning) == "plan"
assert [
(call.function.name, json.loads(call.function.arguments)) for call in calls
] == parse(text)[2]
WEATHER_CALL = call("get_weather", parameter("city", "Paris"))
WEATHER = [("get_weather", {"city": "Paris"})]
def chat_request(tools=STRICT_TOOLS, **body):
return ChatCompletionRequest(
model="beam",
messages=[{"role": "user", "content": "Weather?"}],
tools=tools and [tool.model_dump() for tool in tools],
**body,
)
def named(name):
return {"type": "function", "function": {"name": name}}
def adjust(req, prompt_token_ids=PROMPT_TOKEN_IDS, **kwargs):
make_parser(prompt_token_ids, req.tools, **kwargs).adjust_request(req)
return req
def accepts(req, text):
grammar = Grammar.from_structural_tag(
req.extract_structured_outputs().structural_tag
)
return _is_grammar_accept_string(grammar, text)
@pytest.mark.parametrize(
("tool_choice", "text", "calls"),
[
("required", THOUGHT + WEATHER_CALL, WEATHER),
("required", WEATHER_CALL, None),
(
"required",
THOUGHT
+ call("get_weather", parameter("city", "Paris"), terminator=NEXT_MESSAGE)
+ call("set_count", parameter("count", "7")),
[*WEATHER, ("set_count", {"count": 7})],
),
("required", THOUGHT + call("set_count", parameter("count", "x")), None),
(
"required",
THOUGHT + call("get_weather", parameter("city", "a<|call|>b")),
None,
),
("required", THOUGHT + FINAL + "Paris<|return|>", None),
(named("set_count"), THOUGHT + WEATHER_CALL, None),
(named("get_weather"), THOUGHT + WEATHER_CALL, WEATHER),
(
named("get_weather"),
THOUGHT
+ call("get_weather", parameter("city", "Paris"), terminator=NEXT_MESSAGE)
+ WEATHER_CALL,
None,
),
],
)
def test_forced_tool_calls_follow_the_template_grammar(tool_choice, text, calls):
"""From the message start, the grammar opens with the reasoning, and
accepts only calls the parser reads back."""
req = adjust(chat_request(tool_choice=tool_choice))
assert req._grammar_from_parser
assert accepts(req, text) is (calls is not None)
if calls is not None:
assert parse(text)[2] == calls
def test_reasoning_may_be_skipped_when_thinking_is_disabled():
req = adjust(
chat_request(tool_choice="required"),
chat_template_kwargs={"enable_thinking": False},
)
assert accepts(req, WEATHER_CALL)
assert accepts(req, THOUGHT + WEATHER_CALL)
def test_parallel_tool_calls_false_allows_a_single_call():
req = adjust(chat_request(tool_choice="required", parallel_tool_calls=False))
calls = call("get_weather", parameter("city", "Paris"), terminator=NEXT_MESSAGE)
assert accepts(req, THOUGHT + WEATHER_CALL)
assert not accepts(req, THOUGHT + calls + WEATHER_CALL)
@pytest.mark.parametrize("strict", [False, True])
@pytest.mark.parametrize(
("separator", "accepted"), [("", True), ("\n", True), ("\n\n", False), (" ", False)]
)
def test_parameters_are_separated_by_at_most_a_newline(strict, separator, accepted):
"""A model that skips a parameter's opener must not write whitespace until
`max_tokens`."""
write = _tool(
"write",
{"path": {"type": "string"}, "content": {"type": "string"}},
strict=strict,
)
req = adjust(chat_request([write], tool_choice="required"))
body = parameter("path", "a.py") + separator + parameter("content", "x")
assert accepts(req, THOUGHT + call("write", body)) is accepted
assert accepts(req, THOUGHT + call("write", separator + body)) is accepted
def test_forced_calls_in_prefilled_reasoning_start_after_it():
req = adjust(
chat_request(tool_choice="required"), PROMPT_TOKEN_IDS + [201, 101, 202, 103]
)
assert not req._grammar_from_parser
assert accepts(req, WEATHER_CALL)
def test_forced_calls_cannot_follow_a_prefilled_answer():
with pytest.raises(VLLMValidationError, match="required tool choice"):
adjust(chat_request(tool_choice="required"), PROMPT_TOKEN_IDS + [201, 102, 202])
def test_prefilled_answer_can_only_be_finished():
req = adjust(chat_request(TOOLS), PROMPT_TOKEN_IDS + [201, 102, 202])
assert accepts(req, "Sunny.<|return|>")
assert not accepts(req, "Sunny." + NEXT_MESSAGE + WEATHER_CALL)
def test_forced_call_with_unencodable_tool_name_is_rejected():
tools = [_tool("get weather", {"city": {"type": "string"}})]
with pytest.raises(VLLMValidationError, match="tool name"):
adjust(chat_request(tools, tool_choice="required"))
def test_string_formats_and_patterns_are_held_in_raw_text():
"""Like XGrammar, a strict string keeps its format or pattern, unless a
match could end the value early."""
contact = _tool(
"set_contact",
{
"email": {"type": "string", "format": "email"},
"code": {"type": "string", "pattern": "^[A-Z]{3}$"},
"n": {"type": "integer"},
},
strict=True,
)
contact.function.parameters["required"] = ["email", "code", "n"]
req = adjust(chat_request([contact], tool_choice="required"))
def contact_call(email, code, *rest):
body = parameter("email", email) + parameter("code", code) + "".join(rest)
return THOUGHT + call("set_contact", body)
text = contact_call("a.b@x.co", "ABC", parameter("n", "1"))
calls = make_parser(tools=[contact]).parse(text, req, enable_auto_tools=True)[2]
assert accepts(req, text)
assert normalize_calls(calls) == [
("set_contact", {"email": "a.b@x.co", "code": "ABC", "n": 1})
]
assert not accepts(req, contact_call("ab", "ABC", parameter("n", "1")))
assert not accepts(req, contact_call("a.b@x.co", "ABCD", parameter("n", "1")))
assert not accepts(req, contact_call("a.b@x.co", "ABC"))
# Line breaks could end the value.
dotted = _tool(
"set_code", {"code": {"type": "string", "pattern": "^a.b$"}}, strict=True
)
with pytest.raises(VLLMValidationError, match="required tool choice"):
adjust(chat_request([dotted], tool_choice="required"))
@pytest.mark.parametrize(
("text", "accepted"),
[
(THOUGHT + WEATHER_CALL, True),
(THOUGHT + FINAL + "Sunny.<|return|>", True),
(THOUGHT + FINAL + "Sunny." + NEXT_MESSAGE + WEATHER_CALL, False),
(THOUGHT + call("set_count", parameter("count", "x")), False),
],
)
def test_strict_tools_constrain_automatic_tool_choice(text, accepted):
assert accepts(adjust(chat_request()), text) is accepted
@pytest.mark.parametrize(
("text", "accepted"),
[
(THOUGHT + FINAL + "Sunny.<|return|>", True),
(FINAL + "Sunny.<|return|>", False),
(WEATHER_CALL, False),
(THOUGHT + WEATHER_CALL, False),
(FINAL + "Sunny." + NEXT_MESSAGE + WEATHER_CALL, False),
],
)
def test_tool_choice_none_allows_only_an_answer(text, accepted):
req = adjust(chat_request(TOOLS, tool_choice="none"))
assert req._grammar_from_parser
assert accepts(req, text) is accepted
@pytest.mark.parametrize(("tools", "tool_choice"), [(TOOLS, "auto"), (None, None)])
def test_always_constrained_output_follows_the_template_framing(tools, tool_choice):
req = adjust(chat_request(tools, tool_choice=tool_choice))
count_call = THOUGHT + call("set_count", parameter("count", "3"))
assert accepts(req, THOUGHT + FINAL + "Sunny.<|return|>")
assert not accepts(req, FINAL + "Sunny.<|return|>")
assert not accepts(req, THOUGHT.replace("plan", "") + FINAL + "Sunny.<|return|>")
assert not accepts(req, THOUGHT + "Sunny.")
assert accepts(req, count_call) is (tool_choice == "auto")
@pytest.mark.parametrize("response_format", [None, {"type": "json_object"}])
@pytest.mark.parametrize("strict", [None, True])
@pytest.mark.parametrize("tool_choice", ["auto", "required"])
def test_automatic_choice_lets_calls_name_undeclared_tools(
strict, tool_choice, response_format
):
"""Unless a forced choice limits calls to the declared tools. The parser
passes such calls on, so the harness can tell the model the tool does not
exist."""
bash = _tool("bash", {"command": {"type": "string"}}, strict=strict)
req = adjust(
chat_request([bash], tool_choice=tool_choice, response_format=response_format)
)
undeclared = THOUGHT + call("read_file", parameter("path", "a.py"))
parser = make_parser(tools=[bash])
assert accepts(req, THOUGHT + call("bash", parameter("command", "ls")))
assert accepts(req, undeclared) is (tool_choice == "auto")
calls = parser.parse(undeclared, req, enable_auto_tools=True)[2]
assert normalize_calls(calls) == [("read_file", {"path": "a.py"})]
@pytest.mark.parametrize(
("name", "body", "accepted"),
[
("set_count", parameter("count", "1"), True),
("set_count", parameter("count", "x"), False),
("get_weather", parameter("town", "Paris"), False),
# The grammar cannot enforce `minProperties`.
("find", parameter("x", "1") + parameter("x", "y"), True),
("search_files", parameter("x", "1") + parameter("x", "y"), True),
("set_coun", parameter("x", "1") + parameter("x", "y"), True),
("set_counts", parameter("x", "1") + parameter("x", "y"), True),
],
)
def test_automatic_calls_enforce_declared_arguments_beside_any_other_name(
name, body, accepted
):
find = _tool("find", {"q": {"type": "string"}})
find.function.parameters["minProperties"] = 1
req = adjust(chat_request([*TOOLS, find]))
assert accepts(req, THOUGHT + call(name, body)) is accepted
def test_value_ends_before_an_opener_on_the_next_line():
write = _tool("write", {"path": {"type": "string"}, "content": {"type": "string"}})
req = adjust(chat_request([write], tool_choice="required"))
body = '<rfl:parameter name="content">\nprint(1)\n\n' + parameter("path", "a.py")
text = THOUGHT + call("write", body)
assert accepts(req, text)
calls = make_parser(tools=[write]).parse(text, req, enable_auto_tools=True)[2]
assert normalize_calls(calls) == [
("write", {"content": "print(1)\n", "path": "a.py"})
]
def test_always_constrained_output_can_be_disabled(monkeypatch):
monkeypatch.setenv("VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT", "0")
assert adjust(chat_request(TOOLS)).structured_outputs is None
def test_strict_tools_the_template_cannot_hold_leave_automatic_choice_unconstrained():
dotted = _tool(
"set_code", {"code": {"type": "string", "pattern": "^a.b$"}}, strict=True
)
assert adjust(chat_request([dotted])).structured_outputs is None
def test_generic_parser_does_not_derive_the_template_grammar():
req = chat_request(TOOLS, tool_choice="required")
parser = ResponseTemplateParser(BeamTokenizer(), req.tools, enable_auto_tools=True)
with pytest.raises(VLLMValidationError, match="required tool choice"):
parser.adjust_request(req)
@pytest.mark.parametrize("tools", [None, TOOLS])
def test_structured_output_is_framed_as_the_answer(tools):
schema = {"type": "object", "properties": {"a": {"type": "integer"}}}
req = adjust(
chat_request(
tools,
response_format={
"type": "json_schema",
"json_schema": {"name": "answer", "schema": schema},
},
)
)
assert req.response_format is None
assert accepts(req, THOUGHT + FINAL + '{"a": 1}<|return|>')
assert not accepts(req, THOUGHT + FINAL + "Sunny.<|return|>")
assert accepts(req, THOUGHT + WEATHER_CALL) is (tools is not None)
def test_responses_named_tool_choice_follows_the_template_grammar():
req = adjust(
ResponsesRequest(
input="Weather?",
tools=[
FunctionTool(
type="function",
name="get_weather",
parameters={
"type": "object",
"properties": {"city": {"type": "string"}},
},
strict=True,
)
],
tool_choice=ToolChoiceFunction(type="function", name="get_weather"),
text={"format": {"type": "json_object"}},
)
)
assert req._grammar_from_parser
assert req.text.format.type == "json_object"
assert accepts(req, THOUGHT + WEATHER_CALL)
assert not accepts(req, FINAL + "{}<|return|>")
@pytest.mark.parametrize("status", ["in_progress", "completed"])
@pytest.mark.parametrize("inside_reasoning", [False, True])
def test_responses_preserves_text_config_with_template_grammar(
status, inside_reasoning
):
text = {
"format": {
"type": "json_schema",
"name": "answer",
"strict": True,
"schema": {
"type": "object",
"properties": {"a": {"type": "integer"}},
"required": ["a"],
"additionalProperties": False,
},
},
"verbosity": "low",
}
req = ResponsesRequest(input="Answer", text=text)
expected_text = req.text.model_dump(by_alias=True)
prompt = PROMPT_TOKEN_IDS + ([201, 101, 202] if inside_reasoning else [])
adjust(req, prompt_token_ids=prompt)
params = req.to_sampling_params(default_max_tokens=128)
grammar = Grammar.from_structural_tag(params.structured_outputs.structural_tag)
prefix = "" if inside_reasoning else THOUGHT
assert _is_grammar_accept_string(grammar, prefix + FINAL + '{"a": 1}<|return|>')
assert not _is_grammar_accept_string(grammar, prefix + FINAL + "Sunny.<|return|>")
response = ResponsesResponse.from_request(req, params, "beam", 0, [], status)
assert response.model_dump(by_alias=True)["text"] == expected_text
@pytest.mark.parametrize("reverse_alternatives", [False, True])
def test_reasoning_budget_initializes_from_checkpoint_pattern(
monkeypatch, reverse_alternatives
):
template = deepcopy(BEAM_RESPONSE_TEMPLATE)
alternatives = [
r"^<\|channel\|>analysis<\|message\|>",
r"<\|start\|>assistant<\|channel\|>analysis<\|message\|>",
]
if reverse_alternatives:
alternatives.reverse()
template["fields"]["thinking"] = {
"open_pattern": "(?:" + "|".join(alternatives) + ")",
"close": "<|end|>",
"content": "text",
}
tokenizer = BeamTokenizer(template)
monkeypatch.setattr(
"vllm.config.reasoning.cached_tokenizer_from_config", lambda **_: tokenizer
)
config = ReasoningConfig(reasoning_parser="beam")
config.initialize_token_ids(SimpleNamespace())
assert config.enabled
# The marker must work both after the prompt's assistant header and in
# subsequent complete assistant messages, regardless of regex branch order.
assert config.reasoning_start_token_ids == tokenizer.encode(
"<|channel|>analysis<|message|>"
)
assert config.reasoning_end_token_ids == tokenizer.encode("<|end|>")
assert config.natural_reasoning_end_token_ids == tokenizer.encode("<|end|>")
@pytest.mark.parametrize("pattern", [r"analysis.*", r"(?:analysis|thinking)"])
def test_reasoning_budget_does_not_guess_a_marker_for_other_patterns(pattern):
template = deepcopy(BEAM_RESPONSE_TEMPLATE)
thinking = template["fields"]["thinking"]
del thinking["open"]
thinking["open_pattern"] = pattern
assert BeamReasoningParser(BeamTokenizer(template)).reasoning_start_str is None
tests/parser/test_parse.py
+179
View File
@@ -9,6 +9,9 @@ import pytest
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.parser.abstract_parser import DelegatingParser
from vllm.parser.engine.parser_engine import ParserEngine
from vllm.parser.parser_manager import ParserManager
from vllm.parser.qwen3 import qwen3_config
from vllm.parser.utils import count_history_tool_calls
from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
@@ -417,3 +420,179 @@ def test_parse_auto_tools_no_calls_returns_none(tokenizer):
assert reasoning is None
assert content == PLAIN_TEXT
assert tool_calls is None
# --------------------------------------------------------------------------
# Prompt that already closes the reasoning block
# --------------------------------------------------------------------------
CONTINUE_MESSAGES = [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "Sure,"},
]
REASONED_OUTPUT = "<think>\nlet me think\n</think>\n\n" + PLAIN_TEXT
TOOL_TURN_MESSAGES = [
{"role": "user", "content": "Weather in Paris?"},
{
"role": "assistant",
"content": "<think>\nneed a tool\n</think>\n\n",
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "get_weather", "arguments": "{}"},
}
],
},
{"role": "tool", "tool_call_id": "call_1", "content": "sunny"},
]
def reasoning_parser(tokenizer, name="qwen3"):
parser_cls = ParserManager.get_parser(
reasoning_parser_name=name, tokenizer=tokenizer
)
return parser_cls(tokenizer, None)
def encode(tokenizer, text):
return tokenizer.encode(text, add_special_tokens=False)
def prompt_case(tokenizer, case):
"""(rendered prompt ids, request, model output) for each prompt shape."""
if case == "continue_final_message":
# The template renders an empty, closed think block before "Sure,".
prompt = tokenizer.apply_chat_template(
CONTINUE_MESSAGES,
tokenize=False,
add_generation_prompt=False,
continue_final_message=True,
)
request = make_request(
messages=CONTINUE_MESSAGES,
add_generation_prompt=False,
continue_final_message=True,
)
return encode(tokenizer, prompt), request, " here is the answer"
if case == "template_closed_thinking":
# The template switches thinking off, but the parser is not told
# through chat_template_kwargs (#53284).
prompt = tokenizer.apply_chat_template(
[{"role": "user", "content": "hi"}],
tokenize=False,
add_generation_prompt=True,
enable_thinking=False,
)
return encode(tokenizer, prompt), make_request(), PLAIN_TEXT
if case == "earlier_turn_closed_reasoning":
# Only a previous assistant turn closes reasoning; this one is open.
prompt = tokenizer.apply_chat_template(
TOOL_TURN_MESSAGES,
tools=[{"type": "function", "function": {"name": "get_weather"}}],
tokenize=False,
add_generation_prompt=True,
)
return encode(tokenizer, prompt), make_request(), REASONED_OUTPUT
assert case == "reasoning_open"
prompt = tokenizer.apply_chat_template(
[{"role": "user", "content": "hi"}],
tokenize=False,
add_generation_prompt=True,
)
return encode(tokenizer, prompt), make_request(), REASONED_OUTPUT
@pytest.mark.parametrize("case", ["continue_final_message", "template_closed_thinking"])
def test_parse_prompt_closed_reasoning_is_content(tokenizer, case):
prompt_ids, request, output = prompt_case(tokenizer, case)
parser = reasoning_parser(tokenizer)
parser.set_prompt_token_ids(prompt_ids)
reasoning, content, _ = parser.parse(
output, request, model_output_token_ids=encode(tokenizer, output)
)
assert reasoning is None
assert content == output
assert parser.count_reasoning_tokens(encode(tokenizer, output)) == 0
@pytest.mark.parametrize(
("parser_name", "case"),
[
("qwen3", "reasoning_open"),
("qwen3", "earlier_turn_closed_reasoning"),
# BaseThinkingReasoningParser scans the prompt without turn
# boundaries, so it must not decide non-streaming classification.
("deepseek_r1", "earlier_turn_closed_reasoning"),
],
)
def test_parse_prompt_open_reasoning_still_extracted(tokenizer, parser_name, case):
prompt_ids, request, output = prompt_case(tokenizer, case)
parser = reasoning_parser(tokenizer, parser_name)
parser.set_prompt_token_ids(prompt_ids)
reasoning, content, _ = parser.parse(
output, request, model_output_token_ids=encode(tokenizer, output)
)
assert reasoning.strip() == "let me think"
assert content.strip() == PLAIN_TEXT
assert parser.count_reasoning_tokens(encode(tokenizer, output)) > 0
def test_parse_prompt_ignored_without_turn_boundaries(tokenizer):
"""Without turn boundaries, an earlier turn's </think> looks like this
turn's, so the engine must not start the output in content."""
prompt_ids, request, output = prompt_case(
tokenizer, "earlier_turn_closed_reasoning"
)
parser = ParserEngine(tokenizer, parser_engine_config=qwen3_config(thinking=True))
parser.set_prompt_token_ids(prompt_ids)
reasoning, content, _ = parser.parse(
output, request, model_output_token_ids=encode(tokenizer, output)
)
assert reasoning.strip() == "let me think"
assert content.strip() == PLAIN_TEXT
@pytest.mark.parametrize(
"case",
[
"continue_final_message",
"template_closed_thinking",
"reasoning_open",
"earlier_turn_closed_reasoning",
],
)
def test_prompt_reasoning_streaming_matches_non_streaming(tokenizer, case):
prompt_ids, request, output = prompt_case(tokenizer, case)
output_ids = encode(tokenizer, output)
parser = reasoning_parser(tokenizer)
parser.set_prompt_token_ids(prompt_ids)
reasoning, content, _ = parser.parse(
output, request, model_output_token_ids=output_ids
)
non_streaming = (reasoning or "", content or "")
non_streaming_count = parser.count_reasoning_tokens(output_ids)
parser = reasoning_parser(tokenizer)
streamed_reasoning = streamed_content = ""
for i, token_id in enumerate(output_ids):
delta = parser.parse_delta(
tokenizer.decode([token_id]),
[token_id],
request,
prompt_token_ids=prompt_ids,
finished=i == len(output_ids) - 1,
)
if delta is not None:
streamed_reasoning += delta.reasoning or ""
streamed_content += delta.content or ""
assert (streamed_reasoning, streamed_content) == non_streaming
assert parser.count_reasoning_tokens(output_ids) == non_streaming_count
tests/parser/test_response_template_grammar.py
+492
View File
@@ -0,0 +1,492 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from copy import deepcopy
from typing import Any
import pytest
from xgrammar import Grammar
from xgrammar.testing import _is_grammar_accept_string
from vllm.parser.chat_parsing.response_templates import load_response_template
from vllm.parser.response_template_grammar import (
answer_rest_structural_tag,
tool_call_structural_tag,
)
TEMPLATE: dict[str, Any] = {
"start_anchor": "<assistant>",
"fields": {
"thinking": {"open": "<think>", "close": "</think>", "content": "text"},
"content": {"content": "text"},
"tool_calls": {
"open_pattern": r'<call name="(?P<name>[\w.-]+)">',
"close": "</call>",
"repeats": True,
"content": "xml-inline",
"content_args": {
"tag_pattern": r'<arg name="(?P<key>[^"]+)">(?P<value>.*?)</arg>'
},
"transform": {
"type": "function",
"function": {"name": "{name}", "arguments": "{content}"},
},
},
},
}
ANSWER_TEMPLATE = deepcopy(TEMPLATE)
ANSWER_TEMPLATE["fields"]["content"] = {
"open": "<answer>",
"close": "</answer>",
"content": "text",
}
def tool(parameters, name="f", strict=True) -> dict[str, Any]:
return {"name": name, "parameters": parameters, "strict": strict}
WEATHER = tool(
{
"type": "object",
"properties": {"city": {"type": "string"}, "days": {"type": "integer"}},
"required": ["city"],
},
name="get.weather",
)
PING = {"name": "ping"}
CALL = '<call name="get.weather"><arg name="city">Paris</arg></call>'
def build(
template=TEMPLATE,
functions=(WEATHER, PING),
*,
parallel=True,
thinking=True,
answer=False,
answer_format=None,
sorted_arguments=False,
):
fields = load_response_template(template).fields
return tool_call_structural_tag(
fields["tool_calls"],
list(functions),
parallel_tool_calls=parallel,
thinking=fields["thinking"] if thinking else None,
answer=fields["content"] if answer else None,
answer_format=answer_format,
sorted_arguments=sorted_arguments,
)
def accepts(tag, text):
return _is_grammar_accept_string(Grammar.from_structural_tag(json.dumps(tag)), text)
def call(*arguments):
body = "".join(f'<arg name="{key}">{value}</arg>' for key, value in arguments)
return f'<call name="f">{body}</call>'
@pytest.mark.parametrize(
("text", "accepted"),
[
("<think>plan</think>" + CALL, True),
(CALL, True),
(CALL.removesuffix("</call>"), True),
(
'<call name="get.weather"><arg name="city">Paris</arg>\n'
'<arg name="days">3</arg></call>',
True,
),
(
'<call name="get.weather"><arg name="city">Paris</arg>'
'<arg name="days">soon</arg></call>',
False,
),
('<call name="get.weather"><arg name="days">3</arg></call>', False),
(CALL + '<call name="ping"><arg name="any">x</arg></call>', True),
('<call name="other"></call>', False),
("<think>plan</think>The answer.", False),
],
)
def test_forced_calls_follow_the_template_format(text, accepted):
assert accepts(build(), text) is accepted
def test_single_call_when_parallel_calls_are_disabled():
tag = build(parallel=False)
assert accepts(tag, CALL)
assert not accepts(tag, CALL + CALL)
def test_thinking_is_rejected_when_not_requested():
assert not accepts(build(thinking=False), "<think>plan</think>" + CALL)
@pytest.mark.parametrize(
("text", "accepted"),
[
("<think>plan</think><answer>Sunny.", True),
("<answer>Sunny.</answer>", True),
("<think>plan</think>" + CALL, True),
("<answer>Sunny.</answer>" + CALL, False),
('<call name="get.weather"><arg name="days">soon</arg></call>', False),
],
)
def test_answer_field_lets_the_model_reply_instead(text, accepted):
assert accepts(build(ANSWER_TEMPLATE, answer=True), text) is accepted
@pytest.mark.parametrize("required", [False, True])
@pytest.mark.parametrize(
("text", "optional_only"),
[
("<think>plan</think>" + CALL, False),
("<think>\nplan</think><answer>Sunny.", False),
("<think> \n</think><answer>Sunny.", False),
(CALL, True),
("<answer>Sunny.", True),
("<think></think><answer>Sunny.", True),
("<think><</think><answer>Sunny.", True),
],
)
def test_required_reasoning_opens_the_message_with_nonempty_reasoning(
required, text, optional_only
):
fields = load_response_template(ANSWER_TEMPLATE).fields
tag = tool_call_structural_tag(
fields["tool_calls"],
[WEATHER],
parallel_tool_calls=True,
thinking=fields["thinking"],
answer=fields["content"],
message_start=["<assistant>"],
requires_reasoning=required,
)
assert accepts(tag, text) is (not required or not optional_only)
def test_answer_rest_finishes_the_open_answer():
answer = load_response_template(ANSWER_TEMPLATE).fields["content"]
tag = answer_rest_structural_tag(answer)
assert accepts(tag, "Sunny.</answer>")
assert accepts(tag, "Sunny.")
assert not accepts(tag, "Sunny.</answer>" + CALL)
def test_answer_without_literal_opener_is_not_derived():
assert build(answer=True) is None
def test_answer_holds_the_answer_format():
schema = {"type": "json_schema", "json_schema": {"type": "object"}}
with_tools = build(ANSWER_TEMPLATE, answer=True, answer_format=schema)
answer_only = build(ANSWER_TEMPLATE, (), answer=True, answer_format=schema)
for tag in (with_tools, answer_only):
assert accepts(tag, '<think>plan</think><answer>{"a": 1}</answer>')
assert not accepts(tag, "<answer>Sunny.</answer>")
assert accepts(with_tools, CALL)
assert not accepts(answer_only, CALL)
def test_json_arguments_follow_the_schema():
template = deepcopy(TEMPLATE)
template["fields"]["tool_calls"].update(content="json", content_args={})
tag = build(template, [WEATHER])
assert accepts(tag, '<call name="get.weather">{"city": "Paris"}</call>')
assert not accepts(tag, '<call name="get.weather">{"days": 3}</call>')
END_X = {"tag_pattern": r'<arg name="(?P<key>[^"]+)">(?P<value>.*?)X'}
@pytest.mark.parametrize(
"change",
[
{"content": "json", "content_args": {"unquoted_keys": True}},
{
"transform": {
"type": "function",
"function": {"name": "fixed", "arguments": "{content}"},
}
},
{"content": "kv-lines", "content_args": {}},
# Framing in which argument text can spell a closer, where the parser
# would end the call.
{"close": ["}X", "</call>"], "content_args": END_X},
{
"close": ["}X", "</call>"],
"content_args": {**END_X, "value_parser": {"name": "json"}},
},
{"close": ["7X", "</call>"], "content_args": END_X},
{"close": [">X", "</call>"]},
{"close": ["></call>", "</call>"]},
{"close": ["><a", "</call>"]},
{"close": ["}</call>", "</call>"], "content": "json", "content_args": {}},
],
)
def test_underivable_tool_formats_are_not_derived(change):
template = deepcopy(TEMPLATE)
template["fields"]["tool_calls"].update(change)
assert build(template) is None
def test_non_literal_thinking_delimiters_are_not_derived():
template = deepcopy(TEMPLATE)
del template["fields"]["thinking"]["open"]
template["fields"]["thinking"]["open_pattern"] = r"<think(?: mode=\w+)?>"
assert build(template) is None
assert build(template, thinking=False) is not None
LEVEL = {"type": "object", "properties": {"level": {"$ref": "#/$defs/level"}}}
@pytest.mark.parametrize(
("schema", "value", "accepted"),
[
({"type": "integer", "nullable": True}, "null", True),
({"type": "integer", "nullable": True}, "soon", False),
({"type": ["integer", "null"]}, "null", True),
({"anyOf": [{"type": "integer"}, {"type": "boolean"}]}, "true", True),
({"anyOf": [{"type": "integer"}, {"type": "boolean"}]}, "soon", False),
({"type": ["string", "integer"]}, "soon", True),
({"enum": ["celsius", "fahrenheit"]}, "celsius", True),
({"enum": ["celsius", "fahrenheit"]}, "kelvin", False),
({"type": "string", "const": "on"}, "off", False),
(LEVEL, '{"level": 2}', True),
(LEVEL, '{"level": "high"}', False),
],
)
def test_argument_values_follow_how_the_parser_reads_them(schema, value, accepted):
parameters = {
"type": "object",
"properties": {"v": schema},
"required": ["v"],
"$defs": {"level": {"type": "integer"}},
}
assert accepts(build(functions=[tool(parameters)]), call(("v", value))) is accepted
@pytest.mark.parametrize(
("schema", "strict", "derivable"),
[
# The parser reads values without a type as text.
({"enum": [1, 2]}, True, False),
({"$ref": "#/$defs/level"}, True, False),
({"$ref": "#/$defs/level"}, False, True),
# Raw text cannot hold its closer.
({"enum": ["a</arg>b"]}, True, False),
({"type": "object", "properties": {"x": {"$ref": "#/x"}}}, True, False),
(False, True, False),
# A match could write the start of the value's end.
({"type": "string", "pattern": "^[a-z<]+$"}, True, False),
({"type": "string", "pattern": "^[a-z<]+$"}, False, True),
({"type": "string", "pattern": "^a.b$"}, True, False),
# XGrammar holds length bounds only in JSON.
({"type": "string", "minLength": 3}, True, False),
({"type": "string", "format": "ipv6"}, True, False),
],
)
def test_schemas_text_values_cannot_hold_are_not_derived(schema, strict, derivable):
parameters = {
"type": "object",
"properties": {"v": schema},
"$defs": {"level": {"type": "integer"}},
}
assert (build(functions=[tool(parameters, strict=strict)]) is not None) is derivable
@pytest.mark.parametrize(
("schema", "value", "accepted"),
[
({"type": "string", "format": "email"}, "a.b@x.co", True),
({"type": "string", "format": "email"}, "ab", False),
({"type": "string", "format": "date-time"}, "2026-10-06T12:00:00Z", True),
({"type": "string", "format": "date-time"}, "2026-13-06T12:00:00Z", False),
({"type": "string", "pattern": "^[A-Z]{3}$"}, "ABC", True),
({"type": "string", "pattern": "^[A-Z]{3}$"}, "ABCD", False),
# Like XGrammar, the format overrides the pattern and the length bounds.
(
{"type": "string", "format": "uuid", "pattern": "^x$", "maxLength": 1},
"123e4567-e89b-12d3-a456-426614174000",
True,
),
],
)
def test_text_values_hold_string_formats_and_patterns(schema, value, accepted):
parameters = {"type": "object", "properties": {"v": schema}, "required": ["v"]}
assert accepts(build(functions=[tool(parameters)]), call(("v", value))) is accepted
def test_values_cannot_hold_a_closer():
template = deepcopy(TEMPLATE)
template["fields"]["tool_calls"]["close"] = ["</call>", "STOP"]
text = {"type": "object", "properties": {"v": {"type": "string"}}}
strings = {"type": "array", "items": {"type": "string"}}
json_text = {"type": "object", "properties": {"v": strings}}
assert accepts(build(template, [tool(text)]), call(("v", "ab")))
assert not accepts(build(template, [tool(text)]), call(("v", "aSTOPb")))
assert not accepts(
build(template, [tool(None, strict=False)]), call(("v", "aSTOPb"))
)
assert accepts(build(template, [tool(json_text)]), call(("v", '["ab"]')))
assert not accepts(build(template, [tool(json_text)]), call(("v", '["aSTOPb"]')))
def test_extra_arguments_need_additional_properties():
parameters = {"type": "object", "properties": {"city": {"type": "string"}}}
extra = call(("city", "Paris"), ("zip", "75001"))
assert not accepts(build(functions=[tool(parameters)]), extra)
# Next to declared keys, an extra one could repeat them unchecked.
allowed = {**parameters, "additionalProperties": True}
assert build(functions=[tool(allowed)]) is None
free_form = {"type": "object", "additionalProperties": True}
assert accepts(build(functions=[tool(free_form)]), extra)
def test_only_strict_tools_follow_their_schema():
off_schema = call(("days", "soon"))
parameters = WEATHER["parameters"]
assert accepts(build(functions=[tool(parameters, strict=False)]), off_schema)
assert not accepts(build(functions=[tool(parameters)]), off_schema)
assert accepts(build(functions=[tool(None)]), call())
assert not accepts(build(functions=[tool(None)]), call(("a", "x")))
def test_annotations_keep_strict_arguments_to_their_schema():
parameters = {
**WEATHER["parameters"],
"$schema": "https://json-schema.org/draft/2020-12/schema",
"$id": "weather",
}
tag = build(functions=[tool(parameters)])
assert accepts(tag, call(("city", "Paris"), ("days", "3")))
assert not accepts(tag, call(("city", "Paris"), ("days", "soon")))
@pytest.mark.parametrize(
("arguments", "accepted", "accepted_sorted"),
[
((("path", "a"), ("content", "b")), True, True),
((("content", "b"), ("path", "a")), False, True),
((("path", "a"),), True, True),
((("content", "b"),), False, False),
((("path", "a"), ("path", "c")), False, False),
((("path", "a"), ("mode", "c")), False, False),
],
)
def test_strict_arguments_follow_declared_or_sorted_order(
arguments, accepted, accepted_sorted
):
parameters = {
"properties": {"path": {"type": "string"}, "content": {"type": "string"}},
"required": ["path"],
}
tags = [
build(functions=[tool(parameters)], sorted_arguments=sorted_arguments)
for sorted_arguments in (False, True)
]
assert [accepts(tag, call(*arguments)) for tag in tags] == [
accepted,
accepted_sorted,
]
LOOKAHEAD_TEMPLATE = deepcopy(TEMPLATE)
LOOKAHEAD_TEMPLATE["fields"]["tool_calls"]["content_args"]["tag_pattern"] = (
r'<arg name="(?P<key>[^"]+)">(?P<value>.*?)\n(?:</arg>|(?=<arg name="))'
)
@pytest.mark.parametrize(
("body", "accepted_loose", "accepted_strict"),
[
('<arg name="city">Paris\n</arg>\n<arg name="days">3\n</arg>', True, True),
('<arg name="city">Paris\n<arg name="days">3\n</arg>', True, True),
('<arg name="days">3\n<arg name="city">Paris\n</arg>', True, False),
('<arg name="city">Paris\n<arg name="city">Lyon\n</arg>', True, False),
('<arg name="city">Paris\n<arg name="hours">3\n</arg>', True, False),
('<arg name="city">Paris\n<arg name="', False, False),
('<arg name="city">Paris\n<arg name="days\n</arg>', False, False),
('<arg name="city">a <arg name="b"> c\n</arg>', True, True),
],
)
def test_values_end_where_a_lookahead_end_does(body, accepted_loose, accepted_strict):
parameters = WEATHER["parameters"]
tags = [
build(LOOKAHEAD_TEMPLATE, [tool(parameters, strict=strict)])
for strict in (False, True)
]
text = f'<call name="f">{body}</call>'
assert [accepts(tag, text) for tag in tags] == [accepted_loose, accepted_strict]
@pytest.mark.parametrize("lookahead", [False, True])
@pytest.mark.parametrize(
("keys", "accepted"),
[
("b", True),
("ab", True),
("bc", True),
("bd", True),
("be", True),
("abde", True),
("bcde", True),
("", False),
("a", False),
("acd", False),
("bdc", False),
("bcc", False),
],
)
def test_last_parameter_follows_every_required_one(keys, accepted, lookahead):
parameters = {
"properties": {key: {"type": "string"} for key in "abcde"},
"required": ["b"],
}
tag = build(LOOKAHEAD_TEMPLATE, [tool(parameters)])
ends = ["\n" if lookahead else "\n</arg>"] * (len(keys) - 1) + ["\n</arg>"]
body = "".join(f'<arg name="{k}">x{e}' for k, e in zip(keys, ends))
assert accepts(tag, f'<call name="f">{body}</call>') is accepted
def test_required_argument_without_schema_is_still_required():
parameters = {"type": "object", "properties": {}, "required": ["a"]}
tag = build(functions=[tool(parameters)])
assert accepts(tag, call(("a", "x")))
assert not accepts(tag, call())
assert build(functions=[tool({**parameters, "required": ["a", 1]})]) is None
@pytest.mark.parametrize(
("function", "match"),
[
(tool({}, name="get weather"), "tool name"),
(tool({"properties": {'a"b': {"type": "string"}}}), "parameter name"),
],
)
def test_unencodable_names_are_rejected(function, match):
with pytest.raises(ValueError, match=match):
build(functions=[function])
tests/renderers/test_completions.py
+27
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import io
from collections.abc import Sequence
from dataclasses import dataclass
@@ -376,6 +377,32 @@ class TestRenderPrompt:
assert results[0]["prompt_token_ids"] == tokens
assert results[0]["prompt"] == "[1, 2, 3, 4]"
@pytest.mark.parametrize(
("truncation_side", "expected"),
[("left", [106, 107, 108, 109]), ("right", [100, 101, 102, 103])],
)
def test_detokenization_matches_truncated_tokens(self, truncation_side, expected):
"""Echoed prompt text must match the truncated tokens sent to the engine."""
renderer = _build_renderer(MockModelConfig())
params = TokenizeParams(
max_total_tokens=100,
truncate_prompt_tokens=4,
truncation_side=truncation_side,
needs_detokenization=True,
)
def render():
return renderer.render_prompts(
_preprocess_prompt(renderer.model_config, list(range(100, 110)))
)
results = renderer.tokenize_prompts(render(), params)
async_results = asyncio.run(renderer.tokenize_prompts_async(render(), params))
for result in (*results, *async_results):
assert result["prompt_token_ids"] == expected
assert result["prompt"] == str(expected)
def test_explicit_side_tokenizer_unbounded(self):
renderer = _build_renderer(MockModelConfig())
tests/test_envs.py
+8
View File
@@ -48,6 +48,14 @@ def test_api_key_is_not_compile_factor(monkeypatch: pytest.MonkeyPatch):
assert "VLLM_API_KEY" not in envs.compile_factors()
def test_ray_bundle_indices_is_not_compile_factor(monkeypatch: pytest.MonkeyPatch):
# Engines sharing a placement group get different bundle indices but
# compile the same graphs, so they must share a cache key.
monkeypatch.setenv("VLLM_RAY_BUNDLE_INDICES", "4,5,6,7")
assert "VLLM_RAY_BUNDLE_INDICES" not in envs.compile_factors()
def test_p2p_side_channel_defaults_and_override(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False)
monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False)
tests/v1/attention/test_expand_page_indices.py
+3
View File
@@ -51,6 +51,7 @@ def _run(block_table, cu_num_tokens, kernel_block_size, out_len, device):
block_table.stride(0),
block_table.stride(1),
cu_num_tokens,
cu_num_tokens[1:] - cu_num_tokens[:-1],
KERNEL_BLOCK_SIZE=kernel_block_size,
BLOCK_SIZE=1024,
)
@@ -109,6 +110,7 @@ def test_ragged_batch_leaves_other_rows_untouched():
bt.stride(0),
bt.stride(1),
cu,
cu[1:] - cu[:-1],
KERNEL_BLOCK_SIZE=kbs,
BLOCK_SIZE=1024,
)
@@ -153,6 +155,7 @@ def test_over_provisioned_chunks_do_not_change_result():
bt.stride(0),
bt.stride(1),
cu,
cu[1:] - cu[:-1],
KERNEL_BLOCK_SIZE=kbs,
BLOCK_SIZE=1024,
)
tests/v1/attention/test_mla_backends.py
+144 -7
View File
@@ -10,6 +10,7 @@ Known Issues:
import sys
from types import MethodType, SimpleNamespace
from typing import Literal
import pytest
import torch
@@ -47,6 +48,9 @@ from vllm.v1.attention.backends.mla.prefill.base import MLADimensions
from vllm.v1.attention.backends.mla.prefill.selector import (
MLAPrefillSelectorConfig,
)
from vllm.v1.attention.backends.mla.prefill.trtllm_ragged import (
is_cute_dsl_mla_prefill_supported,
)
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
from vllm.v1.kv_cache_interface import (
@@ -418,6 +422,11 @@ for backend in BACKENDS_TO_TEST:
if query_len_support != QueryLenSupport.SINGLE_ONLY:
SPEC_DECODE_BACKENDS.append(backend)
# Backends whose decode pipeline replays device-trimmed verify batches.
ADAPTIVE_VARLEN_BACKENDS = [
b for b in BACKENDS_TO_TEST if b == AttentionBackendEnum.ROCM_AITER_MLA
]
BACKEND_BLOCK_SIZES = {}
for backend in BACKENDS_TO_TEST:
supported_sizes = backend.get_class().get_supported_kernel_block_sizes()
@@ -470,6 +479,10 @@ BATCH_SPECS = {
"spec_decode_small": BatchSpec(
seq_lens=[128, 256, 512, 1024], query_lens=[4, 4, 4, 4]
),
# Adaptive verification: verify requests trimmed to 1..k+1 tokens on device.
"spec_decode_ragged_small": BatchSpec(
seq_lens=[128, 256, 512, 1024], query_lens=[8, 3, 1, 5]
),
"spec_decode_medium": BatchSpec(
seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[8, 8, 8, 8, 8, 8]
),
@@ -1474,6 +1487,7 @@ def run_attention_backend(
# Force the prefill backend selection (None means auto-select).
vllm_config.attention_config.mla_prefill_backend = prefill_backend
mla_attention_module.backend_supports_prefill_query_quantization.cache_clear()
# Set the current vllm config so that get_current_vllm_config() works
# in the backend implementations
@@ -1581,6 +1595,42 @@ def run_attention_backend(
return output
def _with_even_cpu_query_split(
common_attn_metadata: CommonAttentionMetadata,
) -> CommonAttentionMetadata:
"""Host query lengths as adaptive verification leaves them.
The device keeps the trimmed per-request lengths; the host sees the same
total spread evenly, the first requests taking the remainder
(AdaptiveVerificationManager.compact_batch).
"""
num_reqs = common_attn_metadata.num_reqs
total = int(common_attn_metadata.query_start_loc_cpu[num_reqs])
cpu_lens = torch.full((num_reqs,), total // num_reqs, dtype=torch.int32)
cpu_lens[: total % num_reqs] += 1
query_start_loc_cpu = torch.zeros_like(common_attn_metadata.query_start_loc_cpu)
torch.cumsum(cpu_lens, dim=0, out=query_start_loc_cpu[1 : num_reqs + 1])
return common_attn_metadata.replace(
query_start_loc_cpu=query_start_loc_cpu,
max_query_len=int(
(
common_attn_metadata.query_start_loc[1:]
- common_attn_metadata.query_start_loc[:-1]
)
.max()
.item()
),
)
def _fp8_cache_roundtrip(tensor: torch.Tensor, scale: float) -> torch.Tensor:
"""Reproduce cache storage rounding for the SDPA reference."""
fp8_dtype = current_platform.fp8_dtype()
limits = torch.finfo(fp8_dtype)
quantized = (tensor.float() / scale).clamp(limits.min, limits.max).to(fp8_dtype)
return (quantized.float() * scale).to(tensor.dtype)
def _run_backend_correctness(
default_vllm_config,
dist_init,
@@ -1595,6 +1645,9 @@ def _run_backend_correctness(
qk_nope_head_dim: int,
v_head_dim: int,
chunked_prefill_workspace_size: int | None = None,
adaptive_verification: bool = False,
trtllm_ragged_backend: Literal["trtllm-gen", "cute-dsl"] = "trtllm-gen",
use_prefill_query_quantization: bool = False,
):
"""Test that all backends produce similar outputs to a reference implementation
using torch.nn.functional.scaled_dot_product_attention.
@@ -1626,6 +1679,10 @@ def _run_backend_correctness(
) and AttentionBackendEnum.CUTLASS_MLA in backends_to_test:
# CUTLASS_MLA does not support non-1 Q/K scales
backends_to_test.remove(AttentionBackendEnum.CUTLASS_MLA)
if adaptive_verification:
backends_to_test = [
b for b in backends_to_test if b in ADAPTIVE_VARLEN_BACKENDS
]
if not backends_to_test:
pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}")
@@ -1670,18 +1727,25 @@ def _run_backend_correctness(
hf_config_override["num_attention_heads"]
)
vllm_config.cache_config.cache_dtype = kv_cache_dtype
vllm_config.attention_config.trtllm_ragged_backend = trtllm_ragged_backend
vllm_config.attention_config.use_prefill_query_quantization = (
use_prefill_query_quantization
)
# For spec decode tests, add a speculative_config to set the reorder_batch_threshold
if is_spec_decode_test:
from vllm.config import SpeculativeConfig
# Get the query length from the batch spec (they should all be uniform)
query_len = batch_spec.query_lens[0]
# Get the query length from the batch spec (they should all be uniform,
# except under adaptive verification, where the longest is k+1)
query_len = max(batch_spec.query_lens)
# Set num_speculative_tokens to query_len - 1
# (since threshold is 1 + num_spec_tokens)
# Use ngram method which doesn't require a draft model
vllm_config.speculative_config = SpeculativeConfig(
method="ngram", num_speculative_tokens=query_len - 1
method="ngram",
num_speculative_tokens=query_len - 1,
enable_adaptive_verification=adaptive_verification,
)
device = torch.device(f"{DEVICE_TYPE}:0")
@@ -1756,6 +1820,19 @@ def _run_backend_correctness(
# K_PE (rope component): [s_len, 1, qk_rope_head_dim]
k_pe_full = torch.randn(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
kv_c_cached, k_pe_cached = kv_c_full, k_pe_full
kv_c_prefill, k_pe_prefill = kv_c_full, k_pe_full
if kv_cache_dtype.startswith("fp8"):
kv_c_cached = _fp8_cache_roundtrip(kv_c_full, k_scale)
k_pe_cached = _fp8_cache_roundtrip(k_pe_full, k_scale)
# Prefill reads cached context and unquantized new tokens.
kv_c_prefill = torch.cat(
[kv_c_cached[:context_len], kv_c_full[context_len:]]
)
k_pe_prefill = torch.cat(
[k_pe_cached[:context_len], k_pe_full[context_len:]]
)
# Determine if this sequence uses the decode pipeline or prefill
# pipeline for each backend
# NOTE: For spec decode tests with uniform query_len > 1, backends that
@@ -1805,10 +1882,10 @@ def _run_backend_correctness(
q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
# K: [s_len, kv_lora_rank + qk_rope_head_dim]
# (broadcasted to all heads)
k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)
k_mqa = torch.cat([kv_c_cached, k_pe_cached.squeeze(1)], dim=-1)
k_mqa = k_mqa.unsqueeze(1).expand(-1, num_q_heads, -1)
# V: [s_len, kv_lora_rank] (broadcasted to all heads)
v_mqa = kv_c_full.unsqueeze(1).expand(-1, num_q_heads, -1)
v_mqa = kv_c_cached.unsqueeze(1).expand(-1, num_q_heads, -1)
# Create custom attention mask for decode path:
# - Query tokens can attend to all context tokens
@@ -1837,13 +1914,18 @@ def _run_backend_correctness(
#######################################################
# Prefill path: MHA-style attention with full sequence
# Apply kv_b_proj to the full kv_c tensor
kv_nope_full = torch.einsum("sl,lnh->snh", kv_c_full, kv_b_proj_weight)
kv_nope_full = torch.einsum("sl,lnh->snh", kv_c_prefill, kv_b_proj_weight)
k_nope_full, v_full = kv_nope_full.split([qk_nope_head_dim, v_head_dim], dim=-1)
# Build attention inputs for full sequence
q_mha = torch.cat([q_nope, q_pe], dim=-1) # [q_len, num_heads, total_dim]
k_pe_full_expanded = k_pe_full.expand(-1, num_q_heads, -1)
k_pe_full_expanded = k_pe_prefill.expand(-1, num_q_heads, -1)
k_full = torch.cat([k_nope_full, k_pe_full_expanded], dim=-1)
if use_prefill_query_quantization:
fp8_dtype = current_platform.fp8_dtype()
q_mha = q_mha.to(fp8_dtype).to(dtype)
k_full = k_full.to(fp8_dtype).to(dtype)
v_full = v_full.to(fp8_dtype).to(dtype)
# Create custom attention mask:
# - Query tokens can attend to all context tokens
@@ -1951,6 +2033,13 @@ def _run_backend_correctness(
)
kv_cache_per_block_size[block_size] = kv_cache
if adaptive_verification:
# Only after the cache is filled: prepopulation reads the context
# lengths off the host query lengths.
metadata_per_block_size[block_size] = _with_even_cpu_query_split(
common_attn_metadata
)
# 4. Run vLLM backends and compare
rtol = 1e-2
atol = {
@@ -2111,6 +2200,8 @@ def test_backend_correctness(
_prefill_backend_dimension_params(),
)
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
@pytest.mark.parametrize("trtllm_ragged_backend", ["trtllm-gen", "cute-dsl"])
@pytest.mark.parametrize("use_prefill_query_quantization", [False, True])
def test_chunked_context_backend_correctness(
default_vllm_config,
dist_init,
@@ -2119,8 +2210,26 @@ def test_chunked_context_backend_correctness(
qk_nope_head_dim: int,
v_head_dim: int,
kv_cache_dtype: str,
trtllm_ragged_backend: Literal["trtllm-gen", "cute-dsl"],
use_prefill_query_quantization: bool,
):
"""Split, packed, and context-free requests match the SDPA reference."""
if prefill_backend != MLAPrefillBackendEnum.TRTLLM_RAGGED and (
trtllm_ragged_backend == "cute-dsl" or use_prefill_query_quantization
):
pytest.skip("TRTLLM ragged-specific test configuration")
if use_prefill_query_quantization and kv_cache_dtype == "auto":
pytest.skip("FP8 prefill queries require an FP8 KV cache")
if trtllm_ragged_backend == "cute-dsl":
query_dtype = (
current_platform.fp8_dtype()
if use_prefill_query_quantization
else torch.bfloat16
)
if not is_cute_dsl_mla_prefill_supported(
qk_nope_head_dim + 64, v_head_dim, query_dtype
):
pytest.skip("Unsupported CuTe DSL prefill dimensions or query dtype")
_run_backend_correctness(
default_vllm_config,
dist_init,
@@ -2135,4 +2244,32 @@ def test_chunked_context_backend_correctness(
qk_nope_head_dim=qk_nope_head_dim,
v_head_dim=v_head_dim,
chunked_prefill_workspace_size=1024,
trtllm_ragged_backend=trtllm_ragged_backend,
use_prefill_query_quantization=use_prefill_query_quantization,
)
@pytest.mark.skipif(
not ADAPTIVE_VARLEN_BACKENDS, reason="No adaptive varlen decode backend"
)
def test_adaptive_varlen_decode_correctness(
default_vllm_config, dist_init, workspace_init
):
"""Device query lengths 1..k+1 against an even host split, as adaptive
verification hands them to the target model's attention. TP 8 gives 16
heads, the padded Kimi-K3 count."""
_run_backend_correctness(
default_vllm_config,
dist_init,
workspace_init,
"spec_decode_ragged_small",
"deepseek-ai/DeepSeek-R1",
8,
"fp8",
1.0,
1.0,
None,
128,
128,
adaptive_verification=True,
)
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+32
View File
@@ -63,9 +63,11 @@ class _ExpandPageIndicesKernel:
stride_0,
stride_1,
paged_kv_indptr,
num_real_tokens,
*,
KERNEL_BLOCK_SIZE,
BLOCK_SIZE,
CLAMP_TAIL=False,
):
# Both block-table strides are passed explicitly: the kernel no longer
# assumes the column stride is 1. This stand-in indexes the tensor
@@ -145,6 +147,7 @@ def _builder(
)
),
_mtp_decode_qlen=mtp_decode_qlen,
_varlen_decode_qlen=0,
_uniform_padded_mtp_qo_len=(AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len),
_use_persistent_metadata=False,
kernel_block_size=kernel_block_size,
@@ -500,6 +503,7 @@ def test_non_causal_dcp_block_bypasses_cprr_in_forward(monkeypatch):
impl.qk_rope_head_dim = 64
impl.scale = 576**-0.5
decode = SimpleNamespace(
varlen_q_rows=None,
max_qo_len=4,
qo_indptr=torch.tensor([0, 4], dtype=torch.int32),
paged_kv_indptr=torch.tensor([0, 1], dtype=torch.int32),
@@ -563,6 +567,7 @@ def test_single_token_dcp_decode_returns_unpadded_lse(monkeypatch):
impl.qk_rope_head_dim = 64
impl.scale = head_dim**-0.5
decode = SimpleNamespace(
varlen_q_rows=None,
max_qo_len=1,
qo_indptr=torch.arange(num_tokens + 1, dtype=torch.int32),
paged_kv_indptr=torch.zeros(num_tokens + 1, dtype=torch.int32),
@@ -787,6 +792,7 @@ def test_segmented_dcp_verify_matches_causal_attention(monkeypatch):
torch.tensor([global_seq_len], dtype=torch.int32, device=device),
)
decode = SimpleNamespace(
varlen_q_rows=None,
max_qo_len=qlen,
paged_kv_indptr=torch.tensor(
[0, local_kv.shape[0]], dtype=torch.int32, device=device
@@ -1452,3 +1458,29 @@ def test_a_two_token_bf16_block_is_allowed(monkeypatch):
monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=2, mtp_qlen=8
)
assert metadata.has_persistent_metadata
def _adaptive_config(
*, adaptive: bool = True, dcp: int = 1, num_speculative_tokens: int = 7
) -> SimpleNamespace:
return SimpleNamespace(
speculative_config=SimpleNamespace(
num_speculative_tokens=num_speculative_tokens,
enable_adaptive_verification=adaptive,
parallel_drafting=False,
),
parallel_config=SimpleNamespace(decode_context_parallel_size=dcp),
use_v2_model_runner=True,
)
def test_varlen_cudagraph_bound_follows_adaptive_verification():
bound = AiterMLAMetadataBuilder.get_varlen_cudagraph_max_query_len
assert bound(_adaptive_config(), SimpleNamespace()) == 8
assert bound(_adaptive_config(adaptive=False), SimpleNamespace()) is None
# DCP routes plan per-row windows from host lengths.
assert bound(_adaptive_config(dcp=2), SimpleNamespace()) is None
# Kimi-K3 DSpark draft layers share the target's group; causality is
# decided per batch, so the merged non-causal spec keeps the bound.
shared_draft_group = SimpleNamespace(non_causal_multi_token_decode=True)
assert bound(_adaptive_config(), shared_draft_group) == 8
tests/v1/core/test_kv_cache_config_builder.py
+294
View File
@@ -0,0 +1,294 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for KVCacheConfigBuilder resolution."""
from unittest.mock import MagicMock, patch
import pytest
from vllm.platforms import Platform
from vllm.v1.core.kv_cache_config_builder import (
KVCacheConfigBuilder,
get_kv_cache_config_builder,
)
from vllm.v1.core.kv_cache_planning import (
DefaultKVCacheConfigBuilder,
get_kv_cache_groups_from_builder,
)
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
def _make_vllm_config(supports_grouping: bool = False) -> MagicMock:
"""Create a minimal mock VllmConfig for builder resolution tests.
``supports_grouping`` selects whether the resolved model class declares
``SupportsKVCacheGrouping`` (exposing a ``get_kv_cache_groups`` classmethod).
"""
cfg = MagicMock()
cfg.model_config.supports_kv_cache_grouping = supports_grouping
cfg.attention_config.hisparse_config = None
if supports_grouping:
cfg.model_config.registry.resolve_model_cls.return_value = (
_GroupedModel,
"arch",
)
return cfg
CUSTOM_PATH = "tests.v1.core.test_kv_cache_config_builder.CustomBuilder"
DEFAULT_PATH = "vllm.v1.core.kv_cache_planning.DefaultKVCacheConfigBuilder"
class CustomBuilder(DefaultKVCacheConfigBuilder):
"""A test builder subclass."""
pass
class _GroupedModel:
"""A test model class with a ``get_kv_cache_groups`` classmethod."""
@classmethod
def get_kv_cache_groups(cls, vllm_config, kv_cache_spec):
return []
class _PartialGroupedModel:
"""A test model class whose grouping omits layers."""
@classmethod
def get_kv_cache_groups(cls, vllm_config, kv_cache_spec):
from vllm.v1.kv_cache_interface import KVCacheGroupSpec
return [KVCacheGroupSpec(["a"], kv_cache_spec["a"])]
class TestPlatformHookResolution:
"""The platform hook selects only the placement builder."""
def test_default_hook_returns_default_builder(self):
cfg = _make_vllm_config()
assert Platform.get_kv_cache_config_builder_cls(cfg) == DEFAULT_PATH
class TestBuilderResolution:
@patch("vllm.platforms.current_platform")
@pytest.mark.parametrize(
("builder_cls_path", "expected_cls"),
[
(None, DefaultKVCacheConfigBuilder),
(CUSTOM_PATH, CustomBuilder),
],
)
def test_resolves_hook_selected_builder(
self, mock_platform, builder_cls_path, expected_cls
):
mock_platform.get_kv_cache_config_builder_cls.return_value = (
builder_cls_path or DEFAULT_PATH
)
cfg = _make_vllm_config()
assert type(get_kv_cache_config_builder(cfg)) is expected_cls
@patch("vllm.platforms.current_platform")
def test_resolves_fresh_instance_per_call(self, mock_platform):
mock_platform.get_kv_cache_config_builder_cls.return_value = CUSTOM_PATH
cfg = _make_vllm_config()
assert get_kv_cache_config_builder(cfg) is not get_kv_cache_config_builder(cfg)
def test_public_interface_declares_two_placement_hooks(self):
# get_kv_cache_configs lives in kv_cache_planning and is Core-owned;
# the builder only declares the placement hooks.
assert KVCacheConfigBuilder.__abstractmethods__ == {
"get_kv_cache_groups",
"get_kv_cache_config_from_groups",
}
assert issubclass(DefaultKVCacheConfigBuilder, KVCacheConfigBuilder)
assert (
"pool_bytes_per_block"
not in DefaultKVCacheConfigBuilder.__abstractmethods__
)
assert (
"max_memory_usage_bytes"
not in DefaultKVCacheConfigBuilder.__abstractmethods__
)
class TestModelGroupingPriority:
"""A model class with a get_kv_cache_groups classmethod owns grouping."""
def test_model_classmethod_owns_grouping(self):
cfg = _make_vllm_config(supports_grouping=True)
assert DefaultKVCacheConfigBuilder().get_kv_cache_groups(cfg, {}) == []
def test_model_grouping_rejects_hisparse(self):
cfg = _make_vllm_config(supports_grouping=True)
cfg.attention_config.hisparse_config = object()
with pytest.raises(ValueError, match="HiSparse"):
DefaultKVCacheConfigBuilder().get_kv_cache_groups(cfg, {})
class TestCoreChecksGrouping:
"""Core checks the grouping whether a model or a platform produced it."""
_SPEC = FullAttentionSpec(
block_size=16, num_kv_heads=1, head_size=128, dtype="float16"
)
def test_model_grouping_is_checked(self):
cfg = _make_vllm_config(supports_grouping=True)
cfg.model_config.registry.resolve_model_cls.return_value = (
_PartialGroupedModel,
"arch",
)
with pytest.raises(AssertionError, match="miss layers"):
get_kv_cache_groups_from_builder(
DefaultKVCacheConfigBuilder(), cfg, {"a": self._SPEC, "b": self._SPEC}
)
@pytest.mark.parametrize(
("groups", "match"),
[
([KVCacheGroupSpec(["a"], _SPEC)], "miss layers"),
(
[KVCacheGroupSpec(["a", "b"], _SPEC), KVCacheGroupSpec(["b"], _SPEC)],
"overlaps",
),
],
)
def test_platform_grouping_is_checked(self, groups, match):
class FixedGroupsBuilder(DefaultKVCacheConfigBuilder):
def get_kv_cache_groups(self, vllm_config, kv_cache_spec):
return groups
with pytest.raises(AssertionError, match=match):
get_kv_cache_groups_from_builder(
FixedGroupsBuilder(),
_make_vllm_config(),
{"a": self._SPEC, "b": self._SPEC},
)
class TestPlatformCustomPriority:
"""A vendor platform can override the hook to customize priority."""
def test_platform_builder_wins_over_model_declaration(self):
class PlatformFirstPlatform(Platform):
@classmethod
def get_kv_cache_config_builder_cls(cls, vllm_config):
return CUSTOM_PATH
cfg = _make_vllm_config()
assert PlatformFirstPlatform.get_kv_cache_config_builder_cls(cfg) == CUSTOM_PATH
with patch("vllm.platforms.current_platform", PlatformFirstPlatform):
assert isinstance(get_kv_cache_config_builder(cfg), CustomBuilder)
class TestProjectionPreservesGroupMetadata:
"""PP projection must carry host_resident/role/enable_kv_transfer."""
def test_partial_projection_preserves_metadata(self):
from vllm.v1.core.kv_cache_planning import _project_kv_cache_groups_to_worker
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheGroupRole,
)
spec = FullAttentionSpec(
block_size=16, num_kv_heads=1, head_size=128, dtype="float16"
)
group = MagicMock()
group.layer_names = ["layer1", "layer2", "layer3"]
group.kv_cache_spec = spec
group.is_eagle_group = False
group.host_resident = True
group.enable_kv_transfer = False
group.role = KVCacheGroupRole.HISPARSE_SOURCE
projected = _project_kv_cache_groups_to_worker(
[group], {"layer1": spec, "layer2": spec}
)
assert len(projected) == 1
p = projected[0]
assert p.host_resident is True
assert p.enable_kv_transfer is False
assert p.role is KVCacheGroupRole.HISPARSE_SOURCE
assert p.layer_names == ["layer1", "layer2"]
class TestBytesPerBlockProbe:
"""The single-block probe must not be polluted by num_gpu_blocks_override."""
def test_probe_suppresses_override(self):
import torch
from vllm.config import CacheConfig, ModelConfig, VllmConfig
from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
from vllm.v1.core.kv_cache_planning import (
DefaultKVCacheConfigBuilder,
_pool_bytes_per_block,
_pool_bytes_per_block_from_builder,
)
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheGroupSpec,
)
spec = FullAttentionSpec(
block_size=16, num_kv_heads=1, head_size=128, dtype=torch.float16
)
vc = VllmConfig(
model_config=ModelConfig(max_model_len=100),
cache_config=CacheConfig(block_size=16),
)
resolve_kv_cache_layout(vc, [["BLHNC"]], [spec])
groups = [
KVCacheGroupSpec(["a"], spec),
KVCacheGroupSpec(["b"], spec),
]
pure_c = _pool_bytes_per_block(vc, groups)
vc.cache_config.num_gpu_blocks_override = 500
probed_c = _pool_bytes_per_block_from_builder(
DefaultKVCacheConfigBuilder(), vc, groups
)
assert probed_c == pure_c
assert vc.cache_config.num_gpu_blocks_override == 500
@pytest.mark.parametrize("declares_pool_bytes", [False, True])
def test_several_device_allocations_need_declared_pool_bytes(
self, declares_pool_bytes
):
"""A placement with one allocation per group cannot be sized from one
backing; the builder must declare its per-block bytes."""
from vllm.v1.core.kv_cache_planning import _pool_bytes_per_block_from_builder
from vllm.v1.kv_cache_interface import (
KVCacheConfig,
KVCacheGroupSpec,
KVCacheTensor,
)
class PerGroupAllocationBuilder(DefaultKVCacheConfigBuilder):
def get_kv_cache_config_from_groups(self, vllm_config, groups, num_blocks):
tensors = [
KVCacheTensor(
size=page * num_blocks,
layers=group.layer_names,
layer_stride=page * num_blocks,
block_stride=page,
)
for group, page in zip(groups, (64, 32))
]
return KVCacheConfig(num_blocks, tensors, groups)
def _get_pool_bytes_per_block(self, vllm_config, groups):
return 64 + 32 if declares_pool_bytes else None
vc = _make_vllm_config()
vc.cache_config.num_gpu_blocks_override = None
groups = [KVCacheGroupSpec([name], MagicMock()) for name in "ab"]
builder = PerGroupAllocationBuilder()
if declares_pool_bytes:
assert _pool_bytes_per_block_from_builder(builder, vc, groups) == 96
else:
with pytest.raises(ValueError, match="_get_pool_bytes_per_block"):
_pool_bytes_per_block_from_builder(builder, vc, groups)
tests/v1/distributed/test_async_llm_dp.py
+53 -33
View File
@@ -18,6 +18,7 @@ from vllm.inputs import PromptType
from vllm.outputs import RequestOutput
from vllm.platforms import current_platform
from vllm.sampling_params import RequestOutputKind
from vllm.v1.engine import EngineCoreRequestType
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.v1.engine.core_client import DPLBAsyncMPClient
from vllm.v1.metrics.loggers import StatLoggerBase
@@ -336,12 +337,12 @@ async def _poll_flag(engine: AsyncLLM, want: bool, timeout: float) -> bool:
@pytest.mark.asyncio
async def test_dp_pause_late_request_does_not_block_drain():
"""A request arriving after pause must not leave the coordinator believing
"""A request arriving after pause must not make the coordinator believe
the engines are running.
Paused engines discard START_DP_WAVE, so nothing can report wave
completion afterwards; if forwarding the wake also marks the engines as
running, drain has no way back and can only end in a timeout.
Paused engines start no wave for it, so nothing could report wave
completion afterwards; a running flag set on its behalf would leave drain
no way back but a timeout.
MoE only: wave coordination is enabled iff the model is MoE, so a dense
model never reaches the coordinator state this exercises.
@@ -365,11 +366,9 @@ async def test_dp_pause_late_request_does_not_block_drain():
# A design that never reports the engines as running would satisfy
# every drain assertion below while destroying the signal, so pin it
# down first. The front-end sets its own copy optimistically when it
# forwards the wake, so sample only after several coordinator
# publishes (every 100ms while stats change) have overwritten it.
# Rank 0 reports the wave only after its first step, which can outlast
# the 2s on a cold device, so wait for the report rather than sample it.
# down first. Rank 0 reports the wave only after its first step, which
# can outlast the 2s on a cold device, so wait for the report rather
# than sample it.
await asyncio.sleep(2)
assert not long_request.done(), "the warmup request was too short to sample"
assert await _poll_flag(engine, True, timeout=30), (
@@ -382,24 +381,12 @@ async def test_dp_pause_late_request_does_not_block_drain():
await engine.pause_generation(mode="abort")
await engine.wait_for_requests_to_drain(drain_timeout=30)
# Awaiting add_request guarantees the new-request notification has been
# sent to the coordinator - the message that used to latch the flag.
collector = await engine.add_request(
request_id="late",
prompt=DP_PAUSE_PROMPT,
params=SamplingParams(max_tokens=5),
)
# The front-end marks the engines running off the back of that
# notification. This is what makes the test non-vacuous: it is the
# path that used to leave the coordinator stuck.
assert await _poll_flag(engine, True, timeout=5), (
"the late request did not notify the coordinator"
)
# It must settle back by itself. Unfixed it never does, because the
# paused engines discard the wake and so never report wave completion.
assert await _poll_flag(engine, False, timeout=60)
assert not await _poll_flag(engine, True, timeout=5)
# The late request was held rather than dropped: it completes on resume.
await engine.resume_generation()
@@ -409,6 +396,46 @@ async def test_dp_pause_late_request_does_not_block_drain():
break
@pytest.mark.asyncio
async def test_dp_request_reaching_rank_after_resume_wakes_peers():
"""A request can reach its rank only after resume (e.g. from another API
server), with no other work anywhere. That rank must wake its peers itself,
or it steps alone. Modelled here by holding the request's ADD on the client
side across pause and resume."""
with ExitStack() as after:
engine = AsyncLLM.from_engine_args(
_get_dp_pause_engine_args(expert_parallel=True)
)
after.callback(engine.shutdown)
client = engine.engine_core
params = SamplingParams(max_tokens=5)
await _consume(engine.generate(DP_PAUSE_PROMPT, params, "warmup"))
await engine.pause_generation(mode="abort")
assert await _poll_flag(engine, False, timeout=30)
held: list[tuple] = []
send_input = client._send_input
def hold_add(request_type, request, engine_id=None):
if request_type == EngineCoreRequestType.ADD:
held.append((request_type, request, engine_id))
return asyncio.sleep(0)
return send_input(request_type, request, engine_id)
client._send_input = hold_add
task = asyncio.create_task(
_consume(engine.generate(DP_PAUSE_PROMPT, params, "late"))
)
await asyncio.sleep(2)
assert held
await engine.resume_generation()
client._send_input = send_input
await send_input(*held[0])
await asyncio.wait_for(task, timeout=60)
@pytest.mark.asyncio
async def test_dp_sleep_late_request_does_not_block_drain():
"""The same latch, reached through sleep rather than pause.
@@ -441,14 +468,7 @@ async def test_dp_sleep_late_request_does_not_block_drain():
prompt=DP_PAUSE_PROMPT,
params=SamplingParams(max_tokens=5),
)
assert await _poll_flag(engine, True, timeout=5), (
"the request did not notify the coordinator"
)
# Sleeping engines cannot report wave completion, so a coordinator that
# marked them running when it forwarded the wake never hears otherwise.
assert await _poll_flag(engine, False, timeout=60)
assert not await _poll_flag(engine, True, timeout=5)
await engine.wake_up()
assert not await engine.is_sleeping()
@@ -628,7 +648,7 @@ async def test_dp_pause_keep_race_staggered_engines():
@pytest.mark.asyncio
async def test_dp_pause_barrier_request_deadlock():
"""Test that start_dp_wave is ignored while paused.
"""Test that a request sent while paused does not wake any engine.
Sequence:
1. Pause all engines (PAUSED_ALL).
@@ -697,7 +717,7 @@ async def test_dp_pause_barrier_request_deadlock():
# Yield so generate() preprocessing completes and
# add_request_async is called (which, in buggy code,
# would send FIRST_REQ and wake engine 1).
# would wake engine 1).
for _ in range(200):
await asyncio.sleep(0)
@@ -725,7 +745,7 @@ async def test_dp_pause_barrier_request_deadlock():
for t in mid_barrier_tasks:
t.cancel()
pytest.fail(
"Staggered barrier deadlocked — FIRST_REQ sent while "
"Staggered barrier deadlocked — a request sent while "
"paused caused collective-op mismatch between engines"
)
tests/v1/e2e/general/test_async_scheduling_accuracy.py
+9 -1
View File
@@ -12,6 +12,7 @@ from typing import Any
import pytest
import torch
import vllm.envs as envs
from tests.conftest import HfRunner, VllmRunner
from tests.utils import (
large_gpu_mark,
@@ -168,9 +169,16 @@ def test_with_eagle3_spec_decoding(
reason=("XPU matmul/attention kernels are not batch-invariant"),
)
def test_with_ngram_gpu_spec_decoding(
vllm_runner: type[VllmRunner], hf_runner: type[HfRunner]
vllm_runner: type[VllmRunner],
hf_runner: type[HfRunner],
monkeypatch: pytest.MonkeyPatch,
):
"""Check ngram accuracy and acceptance across scheduling configurations."""
if current_platform.is_rocm():
# hipBLASLt picks GEMM solutions by batch shape, so sync and async runs
# can flip near-tie greedy tokens and shift ngram acceptance.
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True)
ngram_gpu_config = {
"method": "ngram_gpu",
"num_speculative_tokens": 3,
tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py
+22
View File
@@ -8,6 +8,7 @@ import pytest
from tests.models.utils import check_logprobs_close
from vllm import LLM, SamplingParams
from vllm.inputs import TokensPrompt
from vllm.platforms import current_platform
pytestmark = pytest.mark.cpu_model
@@ -112,3 +113,24 @@ def test_prefix_cache_hit_matches_cold_cache(full_prefill_refs):
name_0="cold_cache",
name_1="warm_cache",
)
def test_prefix_cache_hit_reaches_last_full_block():
"""A repeated prompt must hit every full block of the hybrid block size.
The workers raise the block size to fit the GDN state pages, and the
engine core reruns its config checks after its handshake; the raised size
must survive them, or the prefill's state checkpoint misses the block
boundaries and nothing is reused. One chunk holds the whole prompt.
"""
llm = _make_llm(enable_prefix_caching=True, max_num_batched_tokens=2048)
block_size = llm.llm_engine.vllm_config.cache_config.block_size
# Ends mid-block: only the last full block's state can be reused.
num_tokens = block_size + block_size // 2
assert num_tokens <= 2048
prompt = TokensPrompt(prompt_token_ids=[1000 + i for i in range(num_tokens)])
sp = SamplingParams(max_tokens=1)
hits = [llm.generate(prompt, sp)[0].num_cached_tokens for _ in range(2)]
del llm
assert hits == [0, block_size]

Some files were not shown because too many files have changed in this diff Show More