Compare commits
71
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f36d57ec34 | ||
|
|
b8b49e2395 | ||
|
|
da107a59e5 | ||
|
|
ed9a7526b6 | ||
|
|
2427094152 | ||
|
|
fe32e7830b | ||
|
|
afcb580715 | ||
|
|
3f3e2702c2 | ||
|
|
478b49ddec | ||
|
|
cab5c9a2a9 | ||
|
|
774e552397 | ||
|
|
53fa09d085 | ||
|
|
4d93bc35c9 | ||
|
|
586201ebdc | ||
|
|
88f172188b | ||
|
|
880fc032f4 | ||
|
|
6314de8bad | ||
|
|
c91a87f01a | ||
|
|
ea0d045a05 | ||
|
|
0bdfd5eb84 | ||
|
|
0cbc48c4f9 | ||
|
|
2fd0e52252 | ||
|
|
654bd2bca4 | ||
|
|
b623f7ea95 | ||
|
|
0eeba5eec1 | ||
|
|
f69ede495b | ||
|
|
2a2b5ca791 | ||
|
|
689b0eeb9e | ||
|
|
f8e9c56d15 | ||
|
|
e30313220c | ||
|
|
d247a9dc13 | ||
|
|
7c37096620 | ||
|
|
b817b23f7b | ||
|
|
93da882e73 | ||
|
|
0b25cf4419 | ||
|
|
dcdfe66bfa | ||
|
|
68dafcca75 | ||
|
|
1edfd09ffd | ||
|
|
8a9eb40808 | ||
|
|
f91fb2fcf3 | ||
|
|
a045c7425f | ||
|
|
a3a5a5ece5 | ||
|
|
480fadab1b | ||
|
|
279d25f5cb | ||
|
|
54d0c36fff | ||
|
|
9affc17a05 | ||
|
|
816cc73a9b | ||
|
|
2588ec4f0a | ||
|
|
d68f0b220e | ||
|
|
517e74a964 | ||
|
|
48c0d13e65 | ||
|
|
8c3cc98cff | ||
|
|
e4cbc4385d | ||
|
|
6f8b40a23f | ||
|
|
266b9d9c64 | ||
|
|
182c67daf1 | ||
|
|
fd9e91d7e4 | ||
|
|
035733515f | ||
|
|
023808c23d | ||
|
|
985c97a6a8 | ||
|
|
bd0aecdc08 | ||
|
|
8796838910 | ||
|
|
de21863419 | ||
|
|
0910f7e0e1 | ||
|
|
1f6048abe5 | ||
|
|
98f1279815 | ||
|
|
1fd8bd02a4 | ||
|
|
29d69332aa | ||
|
|
4721bb3aa4 | ||
|
|
f46e6be169 | ||
|
|
8b8546da1c |
@@ -16,6 +16,7 @@ steps:
|
||||
- tests/kernels/test_onednn.py
|
||||
- tests/kernels/test_awq_int4_to_int8.py
|
||||
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
|
||||
@@ -24,7 +25,8 @@ steps:
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
|
||||
pytest -x -v -s tests/kernels/test_onednn.py
|
||||
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
|
||||
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
|
||||
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
|
||||
|
||||
- label: CPU-Compatibility Tests
|
||||
depends_on: []
|
||||
|
||||
@@ -37,7 +37,8 @@ function cpu_tests() {
|
||||
pytest -x -v -s tests/kernels/test_onednn.py
|
||||
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
|
||||
pytest -x -v -s tests/kernels/core/test_cpu_activation.py
|
||||
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
|
||||
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
|
||||
|
||||
# skip tests requiring model downloads if HF_TOKEN is not set
|
||||
# due to rate-limits
|
||||
|
||||
Executable
+39
@@ -0,0 +1,39 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
|
||||
|
||||
uv pip install --system -r "${REQUIREMENTS_FILE}"
|
||||
|
||||
NIXL_METADATA=$(python3 - <<'PY'
|
||||
import importlib.metadata as metadata
|
||||
|
||||
import torch
|
||||
|
||||
cuda_version = torch.version.cuda
|
||||
if cuda_version is None:
|
||||
raise SystemExit("torch.version.cuda is not set")
|
||||
|
||||
print(cuda_version.split(".", 1)[0], metadata.version("nixl"))
|
||||
PY
|
||||
)
|
||||
read -r CUDA_MAJOR NIXL_VERSION <<<"${NIXL_METADATA}"
|
||||
|
||||
# nixl>=1.1.0 can install multiple CUDA wheel variants. Keep only the variant
|
||||
# matching this CI image so nixl_ep_cpp links against the available libcudart.
|
||||
uv pip uninstall --system nixl-cu12 nixl-cu13 2>/dev/null || true
|
||||
uv pip install --system --no-deps "nixl-cu${CUDA_MAJOR}==${NIXL_VERSION}"
|
||||
|
||||
python3 - <<'PY'
|
||||
import importlib.metadata as metadata
|
||||
|
||||
for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
|
||||
try:
|
||||
version = metadata.version(package_name)
|
||||
except metadata.PackageNotFoundError:
|
||||
version = "not installed"
|
||||
print(f"{package_name}: {version}")
|
||||
PY
|
||||
+15
-13
@@ -1238,14 +1238,11 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/rpc
|
||||
- tests/entrypoints/serve/instrumentator
|
||||
- tests/tool_use
|
||||
- tests/entrypoints/serve
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/serve/instrumentator
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
|
||||
- pytest -v -s tool_use
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1275,10 +1272,14 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1368,7 +1369,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/tool_parsers
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
|
||||
|
||||
- label: OpenAI API correctness # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2745,14 +2746,11 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/rpc
|
||||
- tests/entrypoints/serve/instrumentator
|
||||
- tests/tool_use
|
||||
- tests/entrypoints/serve
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/serve/instrumentator
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
|
||||
- pytest -v -s tool_use
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2782,10 +2780,14 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
timeout_in_minutes: 180
|
||||
|
||||
@@ -11,7 +11,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
|
||||
key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
|
||||
@@ -22,7 +22,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
@@ -34,7 +34,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
@@ -46,7 +46,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
|
||||
@@ -58,7 +58,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
|
||||
@@ -73,7 +73,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
|
||||
|
||||
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
|
||||
@@ -87,7 +87,7 @@ steps:
|
||||
- vllm/v1/worker/kv_connector_model_runner_mixin.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
|
||||
@@ -102,5 +102,5 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
|
||||
|
||||
@@ -11,7 +11,7 @@ steps:
|
||||
- tests/entrypoints/
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/tool_parsers
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
|
||||
|
||||
- label: Entrypoints Integration (LLM)
|
||||
key: entrypoints-integration-llm
|
||||
@@ -60,9 +60,13 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s tool_use
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -98,14 +102,11 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/rpc
|
||||
- tests/entrypoints/serve/instrumentator
|
||||
- tests/tool_use
|
||||
- tests/entrypoints/serve
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/serve/instrumentator
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
|
||||
- pytest -v -s tool_use
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -153,6 +154,5 @@ steps:
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/model_executor/models/whisper.py
|
||||
commands: # LMEval
|
||||
- pytest -s entrypoints/openai/correctness/
|
||||
|
||||
@@ -86,7 +86,7 @@ steps:
|
||||
- tests/v1/metrics
|
||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# split the test to avoid interference
|
||||
- pytest -v -s -m 'not cpu_test' v1/core
|
||||
|
||||
@@ -45,19 +45,19 @@ steps:
|
||||
- vllm/entrypoints/serve/
|
||||
- vllm/v1/engine/
|
||||
- tests/utils.py
|
||||
# - tests/entrypoints/rpc/test_collective_rpc.py
|
||||
# - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
- tests/entrypoints/serve/disagg/test_serving_tokens.py
|
||||
- tests/entrypoints/serve/instrumentator/test_basic.py
|
||||
- tests/entrypoints/serve/instrumentator/test_metrics.py
|
||||
# - tests/entrypoints/serve/instrumentator/test_sleep.py
|
||||
# - tests/entrypoints/serve/dev/test_sleep.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -v -s entrypoints/rpc/test_collective_rpc.py
|
||||
# - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
|
||||
- pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
|
||||
# - pytest -v -s entrypoints/serve/instrumentator/test_sleep.py
|
||||
# - pytest -v -s entrypoints/serve/dev/test_sleep.py
|
||||
|
||||
- label: Rust Frontend Core Correctness
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -10,7 +10,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Add label
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
with:
|
||||
script: |
|
||||
github.rest.issues.addLabels({
|
||||
|
||||
@@ -14,7 +14,7 @@ jobs:
|
||||
steps:
|
||||
- name: Label issues based on keywords
|
||||
id: label-step
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
with:
|
||||
script: |
|
||||
// Configuration: Add new labels and keywords here
|
||||
@@ -315,7 +315,7 @@ jobs:
|
||||
|
||||
- name: CC users for labeled issues
|
||||
if: steps.label-step.outputs.labels_added != '[]'
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
with:
|
||||
script: |
|
||||
// Configuration: Map labels to GitHub users to CC
|
||||
@@ -392,7 +392,7 @@ jobs:
|
||||
|
||||
- name: Request missing ROCm info from issue author
|
||||
if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug')
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
with:
|
||||
script: |
|
||||
const body = (context.payload.issue.body || '').toLowerCase();
|
||||
|
||||
@@ -12,7 +12,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Update PR description
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
with:
|
||||
script: |
|
||||
const { owner, repo } = context.repo;
|
||||
@@ -55,7 +55,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Post welcome comment for first-time contributors
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
with:
|
||||
script: |
|
||||
const { owner, repo } = context.repo;
|
||||
|
||||
@@ -20,7 +20,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check PR label and author merge count
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
with:
|
||||
script: |
|
||||
const { data: pr } = await github.rest.pulls.get({
|
||||
@@ -49,7 +49,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
|
||||
- uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
|
||||
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
- run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
|
||||
|
||||
@@ -21,7 +21,7 @@ repos:
|
||||
rev: v21.1.2
|
||||
hooks:
|
||||
- id: clang-format
|
||||
exclude: 'csrc/(moe/topk_softmax_kernels.cu|quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
|
||||
exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
|
||||
types_or: [c++, cuda]
|
||||
args: [--style=file, --verbose]
|
||||
- repo: https://github.com/DavidAnson/markdownlint-cli2
|
||||
|
||||
@@ -369,6 +369,18 @@ else()
|
||||
add_compile_definitions(-DVLLM_NUMA_DISABLED)
|
||||
endif()
|
||||
|
||||
# check if the pytorch wheel ships libopenblas.so.
|
||||
set(VLLM_OPENBLAS_LIB "")
|
||||
if (NOT ENABLE_X86_ISA)
|
||||
file(GLOB _VLLM_TORCH_OPENBLAS_LIBS
|
||||
"${TORCH_INSTALL_PREFIX}/lib/libopenblas*.so*")
|
||||
# Note: we don't link openblas directly to _C extension, as it's available through libtorch.so
|
||||
if (_VLLM_TORCH_OPENBLAS_LIBS)
|
||||
list(GET _VLLM_TORCH_OPENBLAS_LIBS 0 VLLM_OPENBLAS_LIB)
|
||||
message(STATUS "CPU OpenBLAS library: ${VLLM_OPENBLAS_LIB}")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
#
|
||||
# Generate CPU attention dispatch header
|
||||
#
|
||||
@@ -387,6 +399,7 @@ endif()
|
||||
#
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/activation.cpp"
|
||||
"csrc/cpu/sgl-kernels/fla.cpp"
|
||||
"csrc/cpu/utils.cpp"
|
||||
"csrc/cpu/spec_decode_utils.cpp"
|
||||
"csrc/cpu/layernorm.cpp"
|
||||
@@ -396,6 +409,13 @@ set(VLLM_EXT_SRC
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp")
|
||||
|
||||
if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
|
||||
VLLM_RVV_VLEN GREATER 0 AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/cpu_wna16.cpp"
|
||||
${VLLM_EXT_SRC})
|
||||
endif()
|
||||
|
||||
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/shm.cpp"
|
||||
@@ -403,6 +423,12 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
|
||||
${VLLM_EXT_SRC})
|
||||
endif()
|
||||
|
||||
if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/shm.cpp"
|
||||
${VLLM_EXT_SRC})
|
||||
endif()
|
||||
|
||||
if(USE_ONEDNN)
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/dnnl_kernels.cpp"
|
||||
@@ -411,7 +437,6 @@ endif()
|
||||
|
||||
if (ENABLE_X86_ISA)
|
||||
set(VLLM_EXT_SRC_SGL
|
||||
"csrc/cpu/sgl-kernels/fla.cpp"
|
||||
"csrc/cpu/sgl-kernels/conv.cpp"
|
||||
"csrc/cpu/sgl-kernels/gemm.cpp"
|
||||
"csrc/cpu/sgl-kernels/gemm_int8.cpp"
|
||||
@@ -423,6 +448,7 @@ if (ENABLE_X86_ISA)
|
||||
"csrc/cpu/sgl-kernels/moe_fp8.cpp")
|
||||
|
||||
set(VLLM_EXT_SRC_AVX512
|
||||
"csrc/cpu/sgl-kernels/fla.cpp"
|
||||
"csrc/cpu/shm.cpp"
|
||||
"csrc/cpu/cpu_wna16.cpp"
|
||||
"csrc/cpu/cpu_fused_moe.cpp"
|
||||
@@ -439,6 +465,7 @@ if (ENABLE_X86_ISA)
|
||||
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
|
||||
|
||||
set(VLLM_EXT_SRC_AVX2
|
||||
"csrc/cpu/sgl-kernels/fla.cpp"
|
||||
"csrc/cpu/utils.cpp"
|
||||
"csrc/cpu/spec_decode_utils.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
@@ -512,6 +539,9 @@ else()
|
||||
USE_SABI 3
|
||||
WITH_SOABI
|
||||
)
|
||||
if (VLLM_OPENBLAS_LIB)
|
||||
target_compile_definitions(_C PRIVATE VLLM_HAS_OPENBLAS)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
message(STATUS "Enabling C extension.")
|
||||
|
||||
@@ -31,7 +31,7 @@ endif()
|
||||
|
||||
if(VLLM_FLASH_ATTN_SRC_DIR)
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn SOURCE_DIR
|
||||
vllm-flash-attn SOURCE_DIR
|
||||
${VLLM_FLASH_ATTN_SRC_DIR}
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
)
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG bce29425653ec0fbc579d329883030e832d15ada
|
||||
GIT_TAG dd62dac706b1cf7895bd99b18c6cb7e7e117ee25
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
@@ -94,6 +94,10 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
|
||||
: reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
|
||||
static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
|
||||
|
||||
explicit FP16Vec16(const c10::Half v)
|
||||
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _f16, LMUL_256)(
|
||||
RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
|
||||
|
||||
explicit FP16Vec16(const FP32Vec16& vec);
|
||||
|
||||
void save(void* ptr) const {
|
||||
@@ -165,6 +169,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
|
||||
|
||||
explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
|
||||
explicit BF16Vec16(const c10::BFloat16 v)
|
||||
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16, LMUL_256)(
|
||||
RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
@@ -290,6 +297,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
}
|
||||
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
|
||||
}
|
||||
explicit BF16Vec16(const c10::BFloat16 v)
|
||||
: reg_fp32(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(static_cast<float>(v),
|
||||
VEC_ELEM_NUM)) {}
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
void save(void* ptr) const {
|
||||
float tmp[16];
|
||||
@@ -629,6 +639,19 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
: reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
|
||||
data.reg, data.reg)) {};
|
||||
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
|
||||
explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
|
||||
const uint64_t q_values = static_cast<uint64_t>(value);
|
||||
auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM);
|
||||
auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM);
|
||||
auto shifts =
|
||||
RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM);
|
||||
auto shifted =
|
||||
RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM);
|
||||
auto idx64 =
|
||||
RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM);
|
||||
auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM);
|
||||
reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM);
|
||||
}
|
||||
explicit FP32Vec16(const FP16Vec16& v);
|
||||
|
||||
#ifdef __riscv_zvfbfmin
|
||||
@@ -641,6 +664,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
|
||||
#endif
|
||||
|
||||
// FP8 stub: dead code on RISC-V (fp8 KV cache is x86-only), needed for
|
||||
// load_b_pair_vec template to compile on all platforms.
|
||||
explicit FP32Vec16(const BF16Vec32&, int) : FP32Vec16() {}
|
||||
|
||||
FP32Vec16 operator+(const FP32Vec16& b) const {
|
||||
return FP32Vec16(
|
||||
RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
|
||||
@@ -891,6 +918,30 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
|
||||
acc = acc.fma(a, b);
|
||||
}
|
||||
|
||||
template <typename VecT>
|
||||
static void interleave_save_16b(const VecT& vec0, const VecT& vec1, void* ptr) {
|
||||
alignas(64) uint16_t values0[VecT::VEC_ELEM_NUM];
|
||||
alignas(64) uint16_t values1[VecT::VEC_ELEM_NUM];
|
||||
vec0.save(values0);
|
||||
vec1.save(values1);
|
||||
|
||||
auto* packed = reinterpret_cast<uint32_t*>(ptr);
|
||||
for (int32_t i = 0; i < VecT::VEC_ELEM_NUM; ++i) {
|
||||
packed[i] = static_cast<uint32_t>(values0[i]) |
|
||||
(static_cast<uint32_t>(values1[i]) << 16);
|
||||
}
|
||||
}
|
||||
|
||||
static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
|
||||
void* ptr) {
|
||||
interleave_save_16b(vec0, vec1, ptr);
|
||||
}
|
||||
|
||||
static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
|
||||
void* ptr) {
|
||||
interleave_save_16b(vec0, vec1, ptr);
|
||||
}
|
||||
|
||||
#ifdef __riscv_zvfbfmin
|
||||
template <>
|
||||
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
|
||||
|
||||
+106
-1
@@ -89,6 +89,35 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
|
||||
}
|
||||
};
|
||||
|
||||
struct FP16Vec16 : public Vec<FP16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
ss16x8x2_t reg;
|
||||
|
||||
explicit FP16Vec16(const void* ptr) {
|
||||
reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
|
||||
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
|
||||
}
|
||||
|
||||
explicit FP16Vec16(bool, const void* ptr) : FP16Vec16(ptr) {}
|
||||
|
||||
explicit FP16Vec16(const FP32Vec16&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
vec_xst(reg.val[0], 0, (signed short*)ptr);
|
||||
vec_xst(reg.val[1], 16, (signed short*)ptr);
|
||||
}
|
||||
|
||||
void save(void* ptr, int elem_num) const {
|
||||
int num = std::max(0, std::min(elem_num, VEC_ELEM_NUM));
|
||||
if (num <= 8) {
|
||||
vec_xst_len(reg.val[0], (signed short*)ptr, num * 2);
|
||||
} else {
|
||||
vec_xst(reg.val[0], 0, (signed short*)ptr);
|
||||
vec_xst_len(reg.val[1], (signed short*)ptr + 8, (num - 8) * 2);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
|
||||
@@ -100,6 +129,8 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
|
||||
}
|
||||
|
||||
explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
|
||||
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
@@ -379,6 +410,8 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
reg.val[3] = vec_xl(48, ptr);
|
||||
}
|
||||
|
||||
explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
|
||||
|
||||
explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
|
||||
|
||||
explicit FP32Vec16(const FP32Vec16& data) {
|
||||
@@ -402,6 +435,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
reg.val[3] = data.reg.val[1];
|
||||
}
|
||||
|
||||
explicit FP32Vec16(const FP16Vec16& v);
|
||||
explicit FP32Vec16(const BF16Vec16& v) {
|
||||
reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
|
||||
reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
|
||||
@@ -735,6 +769,40 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
|
||||
#endif
|
||||
}
|
||||
|
||||
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
|
||||
alignas(16) float temp_fp32[16];
|
||||
alignas(16) c10::Half temp_fp16[16];
|
||||
|
||||
vec_xst(v.reg.val[0], 0, temp_fp32);
|
||||
vec_xst(v.reg.val[1], 16, temp_fp32);
|
||||
vec_xst(v.reg.val[2], 32, temp_fp32);
|
||||
vec_xst(v.reg.val[3], 48, temp_fp32);
|
||||
|
||||
for (int i = 0; i < 16; i++) {
|
||||
temp_fp16[i] = c10::Half(temp_fp32[i]);
|
||||
}
|
||||
|
||||
reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16);
|
||||
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16);
|
||||
}
|
||||
|
||||
inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
|
||||
alignas(16) c10::Half temp_fp16[16];
|
||||
alignas(16) float temp_fp32[16];
|
||||
|
||||
vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16);
|
||||
vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16);
|
||||
|
||||
for (int i = 0; i < 16; i++) {
|
||||
temp_fp32[i] = float(temp_fp16[i]);
|
||||
}
|
||||
|
||||
reg.val[0] = vec_xl(0, temp_fp32);
|
||||
reg.val[1] = vec_xl(16, temp_fp32);
|
||||
reg.val[2] = vec_xl(32, temp_fp32);
|
||||
reg.val[3] = vec_xl(48, temp_fp32);
|
||||
}
|
||||
|
||||
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
|
||||
#ifdef _ARCH_PWR10
|
||||
__vector signed short ret[4];
|
||||
@@ -794,6 +862,43 @@ inline void prefetch(const void* addr) {
|
||||
__asm__ __volatile__("dcbt 0, %0" : : "r"(addr) : "memory");
|
||||
}
|
||||
|
||||
}; // namespace vec_op
|
||||
struct INT8Vec64 {
|
||||
__vector signed char data[4];
|
||||
|
||||
INT8Vec64() = default;
|
||||
|
||||
explicit INT8Vec64(const int8_t* ptr) {
|
||||
data[0] = vec_xl(0, ptr);
|
||||
data[1] = vec_xl(16, ptr);
|
||||
data[2] = vec_xl(32, ptr);
|
||||
data[3] = vec_xl(48, ptr);
|
||||
}
|
||||
|
||||
explicit INT8Vec64(bool, const int8_t* ptr) : INT8Vec64(ptr) {}
|
||||
|
||||
void save(int8_t* ptr) const {
|
||||
vec_xst(data[0], 0, ptr);
|
||||
vec_xst(data[1], 16, ptr);
|
||||
vec_xst(data[2], 32, ptr);
|
||||
vec_xst(data[3], 48, ptr);
|
||||
}
|
||||
|
||||
void save(int8_t* ptr, int elem_num) const {
|
||||
if (elem_num <= 0) return;
|
||||
|
||||
int full_vecs = elem_num / 16;
|
||||
for (int i = 0; i < full_vecs && i < 4; i++) {
|
||||
vec_xst(data[i], i * 16, ptr);
|
||||
}
|
||||
|
||||
int remaining = elem_num % 16;
|
||||
if (remaining > 0 && full_vecs < 4) {
|
||||
vec_xst_len(data[full_vecs], ptr + full_vecs * 16, remaining);
|
||||
}
|
||||
}
|
||||
|
||||
void nt_save(int8_t* ptr) const { save(ptr); }
|
||||
};
|
||||
} // namespace vec_op
|
||||
|
||||
#endif
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#include <ATen/native/CPUBlas.h>
|
||||
|
||||
// Unlike brgemm, PyTorch does not publicly expose at::native::cpublas::gemm
|
||||
// If OpenBLS is available in the PyTorch wheel, we rely on it for fast
|
||||
// bf16:bf16->fp32 GEMMs Otherwise, we fall back to PyTorch reference BLAS path.
|
||||
#if defined(VLLM_HAS_OPENBLAS)
|
||||
extern "C" void sbgemm_(char* transa, char* transb, int* m, int* n, int* k,
|
||||
float* alpha, const at::BFloat16* a, int* lda,
|
||||
const at::BFloat16* b, int* ldb, float* beta, float* c,
|
||||
int* ldc);
|
||||
|
||||
extern "C" void sgemm_(char* transa, char* transb, int* m, int* n, int* k,
|
||||
float* alpha, const float* a, int* lda, const float* b,
|
||||
int* ldb, float* beta, float* c, int* ldc);
|
||||
|
||||
inline char blas_transpose(at::native::TransposeType trans) {
|
||||
switch (trans) {
|
||||
case at::native::TransposeType::NoTranspose:
|
||||
return 'n';
|
||||
case at::native::TransposeType::Transpose:
|
||||
return 't';
|
||||
case at::native::TransposeType::ConjTranspose:
|
||||
return 'c';
|
||||
}
|
||||
return 'n';
|
||||
}
|
||||
|
||||
inline void blas_gemm(at::native::TransposeType transa,
|
||||
at::native::TransposeType transb, int64_t m, int64_t n,
|
||||
int64_t k, float alpha, const at::BFloat16* a,
|
||||
int64_t lda, const at::BFloat16* b, int64_t ldb,
|
||||
float beta, float* c, int64_t ldc) {
|
||||
char transa_ = blas_transpose(transa);
|
||||
char transb_ = blas_transpose(transb);
|
||||
int m_ = static_cast<int>(m);
|
||||
int n_ = static_cast<int>(n);
|
||||
int k_ = static_cast<int>(k);
|
||||
int lda_ = static_cast<int>(lda);
|
||||
int ldb_ = static_cast<int>(ldb);
|
||||
int ldc_ = static_cast<int>(ldc);
|
||||
sbgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
|
||||
c, &ldc_);
|
||||
}
|
||||
|
||||
inline void blas_gemm(at::native::TransposeType transa,
|
||||
at::native::TransposeType transb, int64_t m, int64_t n,
|
||||
int64_t k, float alpha, const float* a, int64_t lda,
|
||||
const float* b, int64_t ldb, float beta, float* c,
|
||||
int64_t ldc) {
|
||||
char transa_ = blas_transpose(transa);
|
||||
char transb_ = blas_transpose(transb);
|
||||
int m_ = static_cast<int>(m);
|
||||
int n_ = static_cast<int>(n);
|
||||
int k_ = static_cast<int>(k);
|
||||
int lda_ = static_cast<int>(lda);
|
||||
int ldb_ = static_cast<int>(ldb);
|
||||
int ldc_ = static_cast<int>(ldc);
|
||||
sgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
|
||||
c, &ldc_);
|
||||
}
|
||||
|
||||
inline void blas_gemm(at::native::TransposeType, at::native::TransposeType,
|
||||
int64_t, int64_t, int64_t, float, const at::Half*,
|
||||
int64_t, const at::Half*, int64_t, float, float*,
|
||||
int64_t) {
|
||||
TORCH_CHECK(false, "CPU OpenBLAS hgemm is not available.");
|
||||
}
|
||||
#else
|
||||
template <typename scalar_t>
|
||||
inline void blas_gemm(at::native::TransposeType transa,
|
||||
at::native::TransposeType transb, int64_t m, int64_t n,
|
||||
int64_t k, float alpha, const scalar_t* a, int64_t lda,
|
||||
const scalar_t* b, int64_t ldb, float beta, float* c,
|
||||
int64_t ldc) {
|
||||
auto gemm = at::native::cpublas::gemm_no_downcast_stub.DEFAULT;
|
||||
gemm(c10::CppTypeToScalarType<scalar_t>::value, transa, transb, m, n, k,
|
||||
at::Scalar(alpha), a, lda, b, ldb, at::Scalar(beta), c, ldc);
|
||||
}
|
||||
#endif
|
||||
+278
-141
@@ -301,25 +301,42 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
// attn = k_beta @ key.transpose(-1, -2)
|
||||
// attn: [B, HV, num_chunk, chunk_size, chunk_size]
|
||||
// transpose and pack for key
|
||||
pack_vnni<scalar_t>(
|
||||
/* dst */ k_transpose,
|
||||
/* src */ curr_k_pad,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ chunk_size);
|
||||
// k_beta @ key.transpose(-1, -2)
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ chunk_size,
|
||||
/* ldc */ chunk_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_k_beta,
|
||||
/* B */ k_transpose,
|
||||
/* C */ curr_attn);
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni<scalar_t>(
|
||||
/* dst */ k_transpose,
|
||||
/* src */ curr_k_pad,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ chunk_size);
|
||||
// k_beta @ key.transpose(-1, -2)
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ chunk_size,
|
||||
/* ldc */ chunk_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_k_beta,
|
||||
/* B */ k_transpose,
|
||||
/* C */ curr_attn);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::Transpose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
chunk_size,
|
||||
chunk_size,
|
||||
qk_head_size,
|
||||
1.0f,
|
||||
curr_k_pad,
|
||||
qk_head_size,
|
||||
curr_k_beta,
|
||||
qk_head_size,
|
||||
0.0f,
|
||||
curr_attn,
|
||||
chunk_size);
|
||||
}
|
||||
// attn = attn * decay_mask
|
||||
for (int64_t m = 0; m < chunk_size; m++) {
|
||||
at::vec::map2<float>(
|
||||
@@ -413,25 +430,42 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
// k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK]
|
||||
// k_cumdecay: [B, HV, num_chunk, chunk_size, EK]
|
||||
// pack for value
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ v_pack,
|
||||
/* src */ curr_v_beta,
|
||||
/* N */ chunk_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
// value = attn @ v_beta
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_attn_reduced,
|
||||
/* B */ v_pack,
|
||||
/* C */ curr_value);
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ v_pack,
|
||||
/* src */ curr_v_beta,
|
||||
/* N */ chunk_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
// value = attn @ v_beta
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_attn_reduced,
|
||||
/* B */ v_pack,
|
||||
/* C */ curr_value);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
v_head_size,
|
||||
chunk_size,
|
||||
chunk_size,
|
||||
1.0f,
|
||||
curr_v_beta,
|
||||
v_head_size,
|
||||
curr_attn_reduced,
|
||||
chunk_size,
|
||||
0.0f,
|
||||
curr_value,
|
||||
v_head_size);
|
||||
}
|
||||
// k_beta_g = k_beta * g.exp().unsqueeze(-1)
|
||||
for (int64_t j = 0; j < chunk_size; j++) {
|
||||
int64_t i = 0;
|
||||
@@ -445,25 +479,42 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
}
|
||||
}
|
||||
// pack for k_beta_g
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ k_beta_g_pack,
|
||||
/* src */ k_beta_g,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ qk_head_size);
|
||||
// k_cumdecay = attn @ k_beta_g
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ qk_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ qk_head_size,
|
||||
/* ldc */ qk_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_attn_reduced,
|
||||
/* B */ k_beta_g_pack,
|
||||
/* C */ k_cumdecay);
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ k_beta_g_pack,
|
||||
/* src */ k_beta_g,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ qk_head_size);
|
||||
// k_cumdecay = attn @ k_beta_g
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ qk_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ qk_head_size,
|
||||
/* ldc */ qk_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_attn_reduced,
|
||||
/* B */ k_beta_g_pack,
|
||||
/* C */ k_cumdecay);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
qk_head_size,
|
||||
chunk_size,
|
||||
chunk_size,
|
||||
1.0f,
|
||||
k_beta_g,
|
||||
qk_head_size,
|
||||
curr_attn_reduced,
|
||||
chunk_size,
|
||||
0.0f,
|
||||
k_cumdecay,
|
||||
qk_head_size);
|
||||
}
|
||||
for (int i = 0; i < chunk_size; i++) {
|
||||
at::vec::map<scalar_t>(
|
||||
[](fVec x) { return x; },
|
||||
@@ -551,25 +602,42 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
|
||||
// attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
|
||||
// k_transpose_i = k_i.transpose(-1, -2)
|
||||
pack_vnni<scalar_t>(
|
||||
/* dst */ k_transpose_i,
|
||||
/* src */ k_i,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ chunk_size);
|
||||
// attn_i = q_i @ k_transpose_i
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ chunk_size,
|
||||
/* ldc */ chunk_size,
|
||||
/* add_C */ false,
|
||||
/* A */ q_i,
|
||||
/* B */ k_transpose_i,
|
||||
/* C */ attn_i);
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni<scalar_t>(
|
||||
/* dst */ k_transpose_i,
|
||||
/* src */ k_i,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ chunk_size);
|
||||
// attn_i = q_i @ k_transpose_i
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ chunk_size,
|
||||
/* ldc */ chunk_size,
|
||||
/* add_C */ false,
|
||||
/* A */ q_i,
|
||||
/* B */ k_transpose_i,
|
||||
/* C */ attn_i);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::Transpose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
chunk_size,
|
||||
chunk_size,
|
||||
qk_head_size,
|
||||
1.0f,
|
||||
k_i,
|
||||
qk_head_size,
|
||||
q_i,
|
||||
qk_head_size,
|
||||
0.0f,
|
||||
attn_i,
|
||||
chunk_size);
|
||||
}
|
||||
// attn_i = attn_i * decay_mask_i
|
||||
for (int64_t m = 0; m < chunk_size; m++) {
|
||||
auto attn_i_m = attn_i + m * chunk_size;
|
||||
@@ -609,28 +677,45 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
}
|
||||
|
||||
// pack for curr_last_recurrent_state
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ curr_last_recurrent_state_pack_reduced,
|
||||
/* src */ curr_last_recurrent_state_reduced,
|
||||
/* N */ qk_head_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ curr_last_recurrent_state_pack_reduced,
|
||||
/* src */ curr_last_recurrent_state_reduced,
|
||||
/* N */ qk_head_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
|
||||
// v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
|
||||
// k_cumdecay_i: [chunk_size, EK]
|
||||
// curr_last_recurrent_state: [EK, EV]
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ k_cumdecay_i_reduced,
|
||||
/* B */ curr_last_recurrent_state_pack_reduced,
|
||||
/* C */ v_prime);
|
||||
// v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
|
||||
// k_cumdecay_i: [chunk_size, EK]
|
||||
// curr_last_recurrent_state: [EK, EV]
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ k_cumdecay_i_reduced,
|
||||
/* B */ curr_last_recurrent_state_pack_reduced,
|
||||
/* C */ v_prime);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
v_head_size,
|
||||
chunk_size,
|
||||
qk_head_size,
|
||||
1.0f,
|
||||
curr_last_recurrent_state_reduced,
|
||||
v_head_size,
|
||||
k_cumdecay_i_reduced,
|
||||
qk_head_size,
|
||||
0.0f,
|
||||
v_prime,
|
||||
v_head_size);
|
||||
}
|
||||
|
||||
// v_new = v_prime = v_i - v_prime
|
||||
// v_i: [chunk_size, EV]
|
||||
@@ -663,41 +748,75 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
}
|
||||
// attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV]
|
||||
// curr_last_recurrent_state: [EK, EV]
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ qg,
|
||||
/* B */ curr_last_recurrent_state_pack_reduced,
|
||||
/* C */ attn_inter);
|
||||
if constexpr (brgemm_supported()) {
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ qg,
|
||||
/* B */ curr_last_recurrent_state_pack_reduced,
|
||||
/* C */ attn_inter);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
v_head_size,
|
||||
chunk_size,
|
||||
qk_head_size,
|
||||
1.0f,
|
||||
curr_last_recurrent_state_reduced,
|
||||
v_head_size,
|
||||
qg,
|
||||
qk_head_size,
|
||||
0.0f,
|
||||
attn_inter,
|
||||
v_head_size);
|
||||
}
|
||||
|
||||
// core_attn_out[:, :, i] = attn_inter + attn_i @ v_new
|
||||
// pack for v_prime
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ v_prime_pack_reduced,
|
||||
/* src */ v_prime_reduced,
|
||||
/* N */ chunk_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
// attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
|
||||
// attn_i: [chunk_size, chunk_size]
|
||||
// v_new: [chunk_size, EV]
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ true,
|
||||
/* A */ attn_i_reduced,
|
||||
/* B */ v_prime_pack_reduced,
|
||||
/* C */ attn_inter);
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ v_prime_pack_reduced,
|
||||
/* src */ v_prime_reduced,
|
||||
/* N */ chunk_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
// attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
|
||||
// attn_i: [chunk_size, chunk_size]
|
||||
// v_new: [chunk_size, EV]
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ true,
|
||||
/* A */ attn_i_reduced,
|
||||
/* B */ v_prime_pack_reduced,
|
||||
/* C */ attn_inter);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
v_head_size,
|
||||
chunk_size,
|
||||
chunk_size,
|
||||
1.0f,
|
||||
v_prime_reduced,
|
||||
v_head_size,
|
||||
attn_i_reduced,
|
||||
chunk_size,
|
||||
1.0f,
|
||||
attn_inter,
|
||||
v_head_size);
|
||||
}
|
||||
|
||||
// core_attn_out[:, :, i] = attn_inter
|
||||
for (int64_t m = 0; m < chunk_size; m++) {
|
||||
@@ -762,17 +881,34 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
/* ld_dst */ chunk_size);
|
||||
// kgv = kg.transpose(-1, -2) @ v_new
|
||||
// v_new: [chunk_size, EV]
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ qk_head_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ kg_transpose,
|
||||
/* B */ v_prime_pack_reduced,
|
||||
/* C */ kgv);
|
||||
if constexpr (brgemm_supported()) {
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ qk_head_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ kg_transpose,
|
||||
/* B */ v_prime_pack_reduced,
|
||||
/* C */ kgv);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
v_head_size,
|
||||
qk_head_size,
|
||||
chunk_size,
|
||||
1.0f,
|
||||
v_prime_reduced,
|
||||
v_head_size,
|
||||
kg_transpose,
|
||||
chunk_size,
|
||||
0.0f,
|
||||
kgv,
|
||||
v_head_size);
|
||||
}
|
||||
// last_recurrent_state = 1) + 2)
|
||||
for (int64_t m = 0; m < qk_head_size; m++) {
|
||||
at::vec::map2<float>(
|
||||
@@ -921,7 +1057,8 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
|
||||
float k_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[k_scale_offset] : 1.0f;
|
||||
int64_t v_offset = si * v_strideS + bi * v_strideB + ni * v_strideH;
|
||||
int64_t o_offset = ((bi * seq_len + si) * v_num_heads + ni) * v_head_dim;
|
||||
float beta_val = 1 / (1 + std::exp(-b_ptr[ni]));
|
||||
// See: https://github.com/sgl-project/sglang/pull/26634
|
||||
float beta_val = 1 / (1 + std::exp(-b_ptr[bi * v_num_heads + ni]));
|
||||
fVec beta_vec = fVec(beta_val);
|
||||
int64_t dvi = 0;
|
||||
for (; dvi <= v_head_dim - VecSize; dvi += VecSize) {
|
||||
|
||||
@@ -4,9 +4,12 @@
|
||||
// clang-format off
|
||||
|
||||
#pragma once
|
||||
#include <ATen/native/CPUBlas.h>
|
||||
|
||||
#include "common.h"
|
||||
#include "blas_gemm.h"
|
||||
|
||||
#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
|
||||
#define CPU_CAPABILITY_AVX512
|
||||
#endif
|
||||
|
||||
// amx-bf16
|
||||
#define TILE_M 16
|
||||
@@ -21,31 +24,39 @@ constexpr int block_size_n() {
|
||||
return 2 * TILE_N;
|
||||
}
|
||||
|
||||
constexpr bool brgemm_supported() {
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
#endif
|
||||
}
|
||||
|
||||
// define threshold using brgemm (intel AMX)
|
||||
template <typename T>
|
||||
inline bool can_use_brgemm(int M);
|
||||
template <>
|
||||
inline bool can_use_brgemm<at::BFloat16>(int M) {
|
||||
return M > 4;
|
||||
return brgemm_supported() && M > 4;
|
||||
}
|
||||
template <>
|
||||
inline bool can_use_brgemm<at::Half>(int M) {
|
||||
return true;
|
||||
return brgemm_supported();
|
||||
}
|
||||
// this requires PyTorch 2.7 or above
|
||||
template <>
|
||||
inline bool can_use_brgemm<int8_t>(int M) {
|
||||
return M > 4;
|
||||
return brgemm_supported() && M > 4;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline bool can_use_brgemm<uint8_t>(int M) {
|
||||
return M > 4;
|
||||
return brgemm_supported() && M > 4;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
|
||||
return M > 4;
|
||||
return brgemm_supported() && M > 4;
|
||||
}
|
||||
|
||||
// work around compiler internal error
|
||||
|
||||
@@ -11,7 +11,9 @@
|
||||
|
||||
#include <ATen/cpu/vec/functional.h>
|
||||
#include <ATen/cpu/vec/vec.h>
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
#include <immintrin.h>
|
||||
#endif
|
||||
namespace {
|
||||
|
||||
using namespace at::vec;
|
||||
|
||||
+10
-8
@@ -5,7 +5,7 @@
|
||||
#include <sys/stat.h>
|
||||
#include <unistd.h>
|
||||
|
||||
#ifdef __aarch64__
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
#include <atomic>
|
||||
#endif
|
||||
|
||||
@@ -38,7 +38,7 @@ struct KernelVecType<c10::Half> {
|
||||
};
|
||||
|
||||
struct ThreadSHMContext {
|
||||
#ifdef __aarch64__
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
// memory model is weaker on AArch64, so we use atomic variables for
|
||||
// consumer (load-acquire) and producer (store-release) to make sure
|
||||
// that a stamp cannot be ready before the corresponding data is ready.
|
||||
@@ -75,7 +75,7 @@ struct ThreadSHMContext {
|
||||
TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
|
||||
TORCH_CHECK((size_t)this % 64 == 0);
|
||||
TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0);
|
||||
#ifdef __aarch64__
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
_curr_thread_stamp[0].store(1, std::memory_order_relaxed);
|
||||
_curr_thread_stamp[1].store(1, std::memory_order_relaxed);
|
||||
_ready_thread_stamp[0].store(0, std::memory_order_relaxed);
|
||||
@@ -124,7 +124,7 @@ struct ThreadSHMContext {
|
||||
}
|
||||
|
||||
char get_curr_stamp(int idx) const {
|
||||
#ifdef __aarch64__
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
return _curr_thread_stamp[idx].load(std::memory_order_acquire);
|
||||
#else
|
||||
return _curr_thread_stamp[idx];
|
||||
@@ -132,7 +132,7 @@ struct ThreadSHMContext {
|
||||
}
|
||||
|
||||
char get_ready_stamp(int idx) const {
|
||||
#ifdef __aarch64__
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
return _ready_thread_stamp[idx].load(std::memory_order_acquire);
|
||||
#else
|
||||
return _ready_thread_stamp[idx];
|
||||
@@ -140,7 +140,7 @@ struct ThreadSHMContext {
|
||||
}
|
||||
|
||||
void next_stamp() {
|
||||
#ifdef __aarch64__
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
_curr_thread_stamp[local_stamp_buffer_idx].fetch_add(
|
||||
1, std::memory_order_release);
|
||||
#else
|
||||
@@ -150,7 +150,7 @@ struct ThreadSHMContext {
|
||||
}
|
||||
|
||||
void commit_ready_stamp() {
|
||||
#ifdef __aarch64__
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
_ready_thread_stamp[local_stamp_buffer_idx].store(
|
||||
_curr_thread_stamp[local_stamp_buffer_idx].load(
|
||||
std::memory_order_relaxed),
|
||||
@@ -186,8 +186,10 @@ struct ThreadSHMContext {
|
||||
break;
|
||||
}
|
||||
++_spinning_count;
|
||||
#ifdef __aarch64__
|
||||
#if defined(__aarch64__)
|
||||
__asm__ __volatile__("yield");
|
||||
#elif defined(__powerpc64__)
|
||||
__asm__ __volatile__("or 1,1,1");
|
||||
#else
|
||||
_mm_pause();
|
||||
#endif // __aarch64__
|
||||
|
||||
+22
-21
@@ -378,7 +378,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
#endif
|
||||
|
||||
// SHM CCL
|
||||
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__))
|
||||
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
|
||||
defined(__powerpc64__)
|
||||
ops.def(
|
||||
"init_shm_manager(str name, int group_size, int rank, int thread_num) -> "
|
||||
"int",
|
||||
@@ -447,6 +448,25 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"bool is_vnni) -> Tensor");
|
||||
ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
|
||||
|
||||
// Adapted from sglang: casual_conv1d kernels
|
||||
ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
|
||||
ops.impl("causal_conv1d_weight_pack", torch::kCPU,
|
||||
&causal_conv1d_weight_pack);
|
||||
ops.def(
|
||||
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
|
||||
"conv_states, Tensor? query_start_loc,"
|
||||
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
|
||||
"int pad_slot_id, bool is_vnni) -> "
|
||||
"Tensor");
|
||||
ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
|
||||
ops.def(
|
||||
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
|
||||
"weight, Tensor? bias, bool silu_activation,"
|
||||
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
|
||||
"bool is_vnni) -> Tensor");
|
||||
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
|
||||
#endif
|
||||
|
||||
// Adapted from sglang: GDN kernels
|
||||
ops.def(
|
||||
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
|
||||
@@ -470,25 +490,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"-> (Tensor, Tensor)");
|
||||
ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
|
||||
|
||||
// Adapted from sglang: casual_conv1d kernels
|
||||
ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
|
||||
ops.impl("causal_conv1d_weight_pack", torch::kCPU,
|
||||
&causal_conv1d_weight_pack);
|
||||
ops.def(
|
||||
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
|
||||
"conv_states, Tensor? query_start_loc,"
|
||||
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
|
||||
"int pad_slot_id, bool is_vnni) -> "
|
||||
"Tensor");
|
||||
ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
|
||||
ops.def(
|
||||
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
|
||||
"weight, Tensor? bias, bool silu_activation,"
|
||||
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
|
||||
"bool is_vnni) -> Tensor");
|
||||
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
|
||||
#endif
|
||||
|
||||
// CPU attention kernels
|
||||
ops.def(
|
||||
"get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
|
||||
@@ -518,7 +519,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
|
||||
|
||||
// WNA16
|
||||
#if defined(__AVX512F__)
|
||||
#if defined(__AVX512F__) || defined(__riscv_v)
|
||||
ops.def(
|
||||
"cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
|
||||
"Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
|
||||
#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp"
|
||||
#include "broadcast_load_epilogue_c2x.hpp"
|
||||
|
||||
/*
|
||||
This file defines custom epilogues for fusing channel scales, token scales,
|
||||
|
||||
@@ -20,7 +20,7 @@
|
||||
|
||||
#include "torch_utils.h"
|
||||
|
||||
#include "../async_util.cuh"
|
||||
#include "async_util.cuh"
|
||||
#include "../cuda_compat.h"
|
||||
#include "../type_convert.cuh"
|
||||
#include "dispatch_utils.h"
|
||||
|
||||
@@ -25,7 +25,7 @@
|
||||
#include <cuda_fp8.h>
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "launch_bounds_utils.h"
|
||||
#include "libtorch_stable/launch_bounds_utils.h"
|
||||
|
||||
// Define before including nvfp4_utils.cuh so the header
|
||||
// can use this macro during compilation.
|
||||
|
||||
@@ -34,7 +34,7 @@
|
||||
static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
|
||||
"MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
|
||||
|
||||
#include "launch_bounds_utils.h"
|
||||
#include "libtorch_stable/launch_bounds_utils.h"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
|
||||
@@ -26,7 +26,7 @@
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "nvfp4_utils.cuh"
|
||||
#include "launch_bounds_utils.h"
|
||||
#include "libtorch_stable/launch_bounds_utils.h"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
|
||||
@@ -26,7 +26,7 @@
|
||||
#include "../../cuda_vec_utils.cuh"
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "launch_bounds_utils.h"
|
||||
#include "libtorch_stable/launch_bounds_utils.h"
|
||||
|
||||
// Define before including nvfp4_utils.cuh so the header
|
||||
// can use this macro during compilation.
|
||||
|
||||
@@ -7,14 +7,11 @@
|
||||
|
||||
#include <torch/csrc/stable/ops.h>
|
||||
|
||||
// NOTE: These headers are intentionally kept in csrc/quantization/gguf/ (not
|
||||
// moved to libtorch_stable) to avoid unnecessary reformatting that would break
|
||||
// git rename detection and pollute blame history.
|
||||
#include "../../../quantization/gguf/ggml-common.h"
|
||||
#include "../../../quantization/gguf/vecdotq.cuh"
|
||||
#include "../../../quantization/gguf/dequantize.cuh"
|
||||
#include "../../../quantization/gguf/mmvq.cuh"
|
||||
#include "../../../quantization/gguf/mmq.cuh"
|
||||
#include "ggml-common.h"
|
||||
#include "vecdotq.cuh"
|
||||
#include "dequantize.cuh"
|
||||
#include "mmvq.cuh"
|
||||
#include "mmq.cuh"
|
||||
#include "moe.cuh"
|
||||
#include "moe_vec.cuh"
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
#include "torch_utils.h"
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#include "../persistent_topk.cuh"
|
||||
#include "persistent_topk.cuh"
|
||||
#endif
|
||||
|
||||
namespace {
|
||||
|
||||
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
|
||||
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
|
||||
ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="v0.1.13"
|
||||
ARG AITER_BRANCH="v0.1.13.post1"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="v1.1.0"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
|
||||
@@ -246,6 +246,12 @@ Every image listed in "image_files" is added to the request in the listed order
|
||||
|
||||
The "image" shorthand accepts the same values as "image_files". The "image_url" field accepts either an OpenAI-style object with a "url" field or a URL string.
|
||||
|
||||
By default, image references are sent to the serving endpoint as provided, with local image paths converted to `file://` URLs.
|
||||
|
||||
If the benchmark client should load local and HTTP(S) images before sending requests, pass `--custom-ensure-client-side-data` to encode them as base64 data URLs on the client side.
|
||||
|
||||
Existing `data:image/...` URLs are already self-contained and are kept unchanged.
|
||||
|
||||
```bash
|
||||
# need a model with vision capability here
|
||||
vllm serve Qwen/Qwen2-VL-7B-Instruct
|
||||
@@ -253,13 +259,13 @@ vllm serve Qwen/Qwen2-VL-7B-Instruct
|
||||
|
||||
```bash
|
||||
# run benchmarking script
|
||||
vllm bench serve--save-result --save-detailed \
|
||||
vllm bench serve --save-result --save-detailed \
|
||||
--backend openai-chat \
|
||||
--model Qwen/Qwen2-VL-7B-Instruct \
|
||||
--endpoint /v1/chat/completions \
|
||||
--dataset-name custom_image \
|
||||
--dataset-path <path-to-your-image-data-jsonl> \
|
||||
--allowed-local-media-path /path/to/image/folder
|
||||
--custom-ensure-client-side-data
|
||||
```
|
||||
|
||||
Note that we need to use the `openai-chat` backend and `/v1/chat/completions` endpoint for multimodal inputs.
|
||||
|
||||
@@ -778,7 +778,7 @@ Then, you can use the OpenAI client as follows:
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
|
||||
video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4"
|
||||
video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
|
||||
|
||||
## Use video url in the payload
|
||||
chat_completion_from_url = client.chat.completions.create(
|
||||
|
||||
@@ -437,6 +437,7 @@ th {
|
||||
| `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |
|
||||
| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ |
|
||||
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
|
||||
| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
|
||||
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
|
||||
| `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
|
||||
|
||||
@@ -100,14 +100,44 @@ For further details on renderer APIs, please refer to [this page](renderer.md).
|
||||
- `/version` - Version information
|
||||
- `/load` - Server load metrics
|
||||
|
||||
## Sleep Mode APIs
|
||||
## Server in development mode
|
||||
|
||||
When using the flag VLLM_SERVER_DEV_MODE=1, you enable development endpoints.
|
||||
|
||||
**SECURITY WARNING: These endpoints should NOT be used in production!**
|
||||
|
||||
### Cache Management APIs
|
||||
|
||||
- `/reset_prefix_cache` - Reset prefix cache (can disrupt service)
|
||||
- `/reset_mm_cache` - Reset multimodal cache (can disrupt service)
|
||||
- `/reset_encoder_cache` - Reset encoder cache (can disrupt service)
|
||||
|
||||
### Weight Transfer APIs (RL Training)
|
||||
|
||||
For further details on Weight Transfer, please refer to [this page](../../training/weight_transfer/README.md).
|
||||
|
||||
- `/pause` - Pause generation (causes denial of service)
|
||||
- `/resume` - Resume generation
|
||||
- `/is_paused` - Check if generation is paused
|
||||
- `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
|
||||
- `/update_weights` - Update model weights (can alter model behavior)
|
||||
- `/get_world_size` - Get distributed world size
|
||||
|
||||
### Collective RPC
|
||||
|
||||
- `/collective_rpc` - Execute arbitrary RPC methods on the engine (extremely dangerous)
|
||||
|
||||
### Server info
|
||||
|
||||
- `/server_info` - Get detailed server configuration
|
||||
|
||||
### Sleep Mode APIs
|
||||
|
||||
For further details on sleep mode, please refer to [this page](../../features/sleep_mode.md).
|
||||
|
||||
- `/sleep` - Put engine to sleep (causes denial of service)
|
||||
- `/wake_up` - Wake engine from sleep
|
||||
- `/is_sleeping` - Check if engine is sleeping
|
||||
- `/collective_rpc` - Execute arbitrary RPC methods on the engine (extremely dangerous)
|
||||
|
||||
## Chat Template
|
||||
|
||||
|
||||
@@ -84,7 +84,10 @@ Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightT
|
||||
|
||||
## Receiving Weights (Inference Side)
|
||||
|
||||
The inference side triggers weight reception using the four-phase protocol — `init_weight_transfer_engine`, `start_weight_update`, `update_weights`, `finish_weight_update`. The init phase is shown [above](#initialization). The remaining three steps are:
|
||||
The inference side triggers weight reception using the four-phase protocol:
|
||||
`init_weight_transfer_engine`, `start_weight_update`, `update_weights`,
|
||||
`finish_weight_update`. The init phase is shown [above](#initialization). The
|
||||
remaining three steps are:
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
|
||||
@@ -108,12 +111,24 @@ llm.update_weights(
|
||||
llm.finish_weight_update()
|
||||
```
|
||||
|
||||
The `names`, `dtype_names`, and `shapes` lists describe each parameter. These must match the order in which the trainer iterates over its parameters.
|
||||
The `names`, `dtype_names`, and `shapes` lists describe each parameter. These
|
||||
must match the order in which the trainer iterates over its parameters.
|
||||
|
||||
`start_weight_update` must be called before `update_weights`, and `finish_weight_update` must be called after all weight chunks have been transferred. The `is_checkpoint_format` flag controls whether layerwise reload processing is applied (`True` for checkpoint-format weights, `False` for pre-processed kernel-format weights).
|
||||
`start_weight_update` must be called before `update_weights`, and
|
||||
`finish_weight_update` must be called after all weight chunks have been
|
||||
transferred. The `is_checkpoint_format` flag controls whether layerwise reload
|
||||
processing is applied (`True` for checkpoint-format weights, `False` for
|
||||
pre-processed kernel-format weights).
|
||||
|
||||
Sparse NCCL patches still use `update_kind="sparse_flat"` inside
|
||||
`update_info`, but they should be wrapped in
|
||||
`start_weight_update(is_checkpoint_format=False)` because sparse patches apply
|
||||
directly to runtime/kernel-format parameters. The current sparse MVP requires
|
||||
`TP=1` and `PP=1`.
|
||||
|
||||
## Examples
|
||||
|
||||
- [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
|
||||
- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1`
|
||||
- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
|
||||
- [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
|
||||
|
||||
@@ -203,7 +203,7 @@ def run_multi_image(model: str, max_completion_tokens: int) -> None:
|
||||
|
||||
# Video input inference
|
||||
def run_video(model: str, max_completion_tokens: int) -> None:
|
||||
video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4"
|
||||
video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
|
||||
video_base64 = encode_base64_content_from_url(video_url)
|
||||
|
||||
## Use video url in the payload
|
||||
|
||||
@@ -0,0 +1,526 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Demonstrates dense-vs-sparse NCCL weight syncing with a real model.
|
||||
|
||||
This example mirrors the validation story used for the sparse NCCL MVP:
|
||||
both the dense update path and the sparse patch path start from the same real
|
||||
checkpoint and apply the same deterministic trainer-side patch. The script then
|
||||
checks that greedy 1-token outputs match between the dense and sparse vLLM
|
||||
engines after the update.
|
||||
|
||||
The example performs the following steps:
|
||||
* Load a training model on one GPU via a Ray actor.
|
||||
* Launch a vLLM engine with the same real model on a second GPU.
|
||||
* Verify trainer vs vLLM baseline agreement before any update.
|
||||
* Apply a deterministic patch to ``model.embed_tokens.weight`` on the trainer.
|
||||
* Run a dense NCCL update into a fresh vLLM engine and collect post-update
|
||||
outputs.
|
||||
* Reset the trainer back to the baseline checkpoint.
|
||||
* Apply the same deterministic patch again.
|
||||
* Run a sparse NCCL update into another fresh vLLM engine and collect
|
||||
post-update outputs.
|
||||
* Compare dense vs sparse baseline outputs, dense vs sparse post-update
|
||||
outputs, estimated payload sizes, and trainer-side send times.
|
||||
|
||||
Current sparse weight transfer MVP limitations:
|
||||
* ``TP=1`` and ``PP=1`` only
|
||||
* sparse updates use runtime/kernel-format parameter names
|
||||
* sparse updates are not composable with checkpoint-format or packed updates
|
||||
|
||||
This example assumes a single-node cluster with two GPUs.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
|
||||
import ray
|
||||
import torch
|
||||
from ray.util.placement_group import placement_group
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import WeightTransferConfig
|
||||
from vllm.distributed.weight_transfer.base import SparseWeightPatch
|
||||
from vllm.distributed.weight_transfer.nccl_engine import (
|
||||
NCCLTrainerSendWeightsArgs,
|
||||
NCCLWeightTransferEngine,
|
||||
)
|
||||
from vllm.utils.network_utils import get_ip, get_open_port
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
PATCHED_PARAM_NAME = "model.embed_tokens.weight"
|
||||
MAX_PATCH_ROWS = 32
|
||||
PROMPTS = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
SAMPLING_PARAMS = SamplingParams(temperature=0.0, max_tokens=1)
|
||||
|
||||
|
||||
class MyLLM(LLM):
|
||||
"""Configure the vLLM worker for Ray placement group execution."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
class TrainModel:
|
||||
"""Ray actor that owns the trainer-side model and deterministic patch state."""
|
||||
|
||||
def __init__(self, model_name: str):
|
||||
self.model_name = model_name
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
if self.tokenizer.pad_token_id is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
self.model = None
|
||||
self.patched_param = None
|
||||
self.pending_sparse_patches: list[SparseWeightPatch] | None = None
|
||||
self.model_update_group = None
|
||||
self.master_address = get_ip()
|
||||
self.port = get_open_port()
|
||||
self.reset_model()
|
||||
|
||||
def reset_model(self) -> None:
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
self.model_name,
|
||||
torch_dtype=torch.bfloat16,
|
||||
).to("cuda:0")
|
||||
self.model.eval()
|
||||
|
||||
try:
|
||||
self.patched_param = self.model.get_parameter(PATCHED_PARAM_NAME)
|
||||
except AttributeError as exc:
|
||||
raise RuntimeError(
|
||||
f"Expected trainer model to expose `{PATCHED_PARAM_NAME}`"
|
||||
) from exc
|
||||
|
||||
self.pending_sparse_patches = None
|
||||
|
||||
def create_rendezvous(self) -> tuple[str, int]:
|
||||
self.port = get_open_port()
|
||||
return self.master_address, self.port
|
||||
|
||||
def init_weight_transfer_group(self, world_size: int) -> None:
|
||||
self.model_update_group = NCCLWeightTransferEngine.trainer_init(
|
||||
dict(
|
||||
master_address=self.master_address,
|
||||
master_port=self.port,
|
||||
world_size=world_size,
|
||||
)
|
||||
)
|
||||
|
||||
def get_dense_update_info(self, packed: bool = False) -> tuple[dict, int]:
|
||||
names = []
|
||||
dtype_names = []
|
||||
shapes = []
|
||||
payload_bytes = 0
|
||||
for name, param in self.model.named_parameters():
|
||||
names.append(name)
|
||||
dtype_names.append(str(param.dtype).split(".")[-1])
|
||||
shapes.append(list(param.shape))
|
||||
payload_bytes += param.numel() * param.element_size()
|
||||
|
||||
return (
|
||||
dict(
|
||||
names=names,
|
||||
dtype_names=dtype_names,
|
||||
shapes=shapes,
|
||||
packed=packed,
|
||||
),
|
||||
payload_bytes,
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(
|
||||
self,
|
||||
prompts: Sequence[str],
|
||||
max_new_tokens: int = 1,
|
||||
) -> list[dict[str, object]]:
|
||||
generations = []
|
||||
for prompt in prompts:
|
||||
model_inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
|
||||
output = self.model.generate(
|
||||
**model_inputs,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
pad_token_id=self.tokenizer.pad_token_id,
|
||||
)
|
||||
new_token_ids = output[0, model_inputs["input_ids"].shape[1] :].tolist()
|
||||
generations.append(
|
||||
{
|
||||
"token_ids": new_token_ids,
|
||||
"text": self.tokenizer.decode(
|
||||
new_token_ids,
|
||||
skip_special_tokens=False,
|
||||
),
|
||||
}
|
||||
)
|
||||
return generations
|
||||
|
||||
def prepare_sparse_patch(
|
||||
self,
|
||||
prompts: Sequence[str],
|
||||
max_patch_rows: int = MAX_PATCH_ROWS,
|
||||
) -> tuple[dict[str, object], list[int], str, int]:
|
||||
selected_token_ids: list[int] = []
|
||||
special_ids = set(self.tokenizer.all_special_ids)
|
||||
for prompt in prompts:
|
||||
token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"]
|
||||
for token_id in token_ids:
|
||||
if token_id in special_ids or token_id in selected_token_ids:
|
||||
continue
|
||||
selected_token_ids.append(token_id)
|
||||
if len(selected_token_ids) == max_patch_rows:
|
||||
break
|
||||
if len(selected_token_ids) == max_patch_rows:
|
||||
break
|
||||
|
||||
if not selected_token_ids:
|
||||
raise ValueError("Could not derive any non-special token IDs to patch")
|
||||
|
||||
vocab_size = self.patched_param.shape[0]
|
||||
next_token_id = selected_token_ids[-1]
|
||||
while len(selected_token_ids) < max_patch_rows:
|
||||
next_token_id = (next_token_id + 1) % vocab_size
|
||||
if next_token_id in special_ids or next_token_id in selected_token_ids:
|
||||
continue
|
||||
selected_token_ids.append(next_token_id)
|
||||
|
||||
row_ids = torch.tensor(
|
||||
selected_token_ids,
|
||||
device=self.patched_param.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
hidden_size = self.patched_param.shape[1]
|
||||
column_offsets = torch.arange(
|
||||
hidden_size,
|
||||
device=self.patched_param.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
# Rotate the selected embedding rows instead of zeroing them so the
|
||||
# patch remains deterministic while avoiding a degenerate collapse
|
||||
# to the same special token after the update.
|
||||
replacement_rows = self.patched_param[row_ids].roll(shifts=1, dims=0)
|
||||
self.patched_param[row_ids] = replacement_rows
|
||||
|
||||
flat_indices = (
|
||||
row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1)
|
||||
)
|
||||
flat_values = self.patched_param[row_ids].reshape(-1).contiguous()
|
||||
self.pending_sparse_patches = [
|
||||
SparseWeightPatch(
|
||||
name=PATCHED_PARAM_NAME,
|
||||
indices=flat_indices.to(torch.int32),
|
||||
values=flat_values,
|
||||
)
|
||||
]
|
||||
patch_digest = hashlib.sha256(
|
||||
self.pending_sparse_patches[0].indices.cpu().numpy().tobytes()
|
||||
+ self.pending_sparse_patches[0]
|
||||
.values.detach()
|
||||
.float()
|
||||
.cpu()
|
||||
.numpy()
|
||||
.tobytes()
|
||||
).hexdigest()
|
||||
|
||||
sparse_payload_bytes = (
|
||||
flat_indices.numel() * torch.tensor([], dtype=torch.int32).element_size()
|
||||
+ flat_values.numel() * flat_values.element_size()
|
||||
)
|
||||
update_info = dict(
|
||||
names=[PATCHED_PARAM_NAME],
|
||||
dtype_names=[str(self.patched_param.dtype).split(".")[-1]],
|
||||
shapes=[list(self.patched_param.shape)],
|
||||
num_updates_list=[flat_indices.numel()],
|
||||
update_kind="sparse_flat",
|
||||
)
|
||||
return update_info, selected_token_ids, patch_digest, sparse_payload_bytes
|
||||
|
||||
def broadcast_weights(self, packed: bool = False) -> float:
|
||||
if self.model_update_group is None:
|
||||
raise RuntimeError("Weight transfer group is not initialized")
|
||||
|
||||
trainer_args = NCCLTrainerSendWeightsArgs(
|
||||
group=self.model_update_group,
|
||||
packed=packed,
|
||||
)
|
||||
start = time.perf_counter()
|
||||
NCCLWeightTransferEngine.trainer_send_weights(
|
||||
iterator=self.model.named_parameters(),
|
||||
trainer_args=trainer_args,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
return (time.perf_counter() - start) * 1000.0
|
||||
|
||||
def broadcast_pending_sparse_patch(self) -> float:
|
||||
if self.model_update_group is None:
|
||||
raise RuntimeError("Weight transfer group is not initialized")
|
||||
if self.pending_sparse_patches is None:
|
||||
raise RuntimeError("Sparse patch has not been prepared")
|
||||
|
||||
start = time.perf_counter()
|
||||
NCCLWeightTransferEngine.trainer_send_sparse_weights(
|
||||
iter(self.pending_sparse_patches),
|
||||
NCCLTrainerSendWeightsArgs(group=self.model_update_group),
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
self.pending_sparse_patches = None
|
||||
return (time.perf_counter() - start) * 1000.0
|
||||
|
||||
|
||||
def launch_llm(
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
):
|
||||
return ray.remote(
|
||||
num_cpus=0,
|
||||
num_gpus=0,
|
||||
scheduling_strategy=scheduling_inference,
|
||||
)(MyLLM).remote(
|
||||
model=MODEL_NAME,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=1,
|
||||
distributed_executor_backend="ray",
|
||||
gpu_memory_utilization=0.7,
|
||||
weight_transfer_config=WeightTransferConfig(backend="nccl"),
|
||||
)
|
||||
|
||||
|
||||
def collect_vllm_generations(llm_handle) -> list[dict[str, object]]:
|
||||
outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS))
|
||||
generations = []
|
||||
for output in outputs:
|
||||
generations.append(
|
||||
{
|
||||
"token_ids": output.outputs[0].token_ids,
|
||||
"text": output.outputs[0].text,
|
||||
}
|
||||
)
|
||||
return generations
|
||||
|
||||
|
||||
def token_sequences_match(
|
||||
left: Sequence[dict[str, object]],
|
||||
right: Sequence[dict[str, object]],
|
||||
) -> bool:
|
||||
return [item["token_ids"] for item in left] == [item["token_ids"] for item in right]
|
||||
|
||||
|
||||
def print_generations(label: str, prompts: Sequence[str], generations) -> None:
|
||||
print(f"\n{label}")
|
||||
print("-" * 50)
|
||||
for prompt, generation in zip(prompts, generations):
|
||||
print(f"Prompt: {prompt!r}")
|
||||
print(f"Token IDs: {generation['token_ids']}")
|
||||
print(f"Text: {generation['text']!r}")
|
||||
print("-" * 50)
|
||||
|
||||
|
||||
def run_dense_phase(
|
||||
train_model,
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
) -> dict[str, object]:
|
||||
ray.get(train_model.reset_model.remote())
|
||||
llm = launch_llm(scheduling_inference)
|
||||
try:
|
||||
dense_before = collect_vllm_generations(llm)
|
||||
|
||||
ray.get(llm.sleep.remote(level=0))
|
||||
master_address, master_port = ray.get(train_model.create_rendezvous.remote())
|
||||
world_size = ray.get(llm.get_world_size.remote()) + 1
|
||||
inference_init = llm.init_weight_transfer_engine.remote(
|
||||
dict(
|
||||
init_info=dict(
|
||||
master_address=master_address,
|
||||
master_port=master_port,
|
||||
rank_offset=1,
|
||||
world_size=world_size,
|
||||
)
|
||||
)
|
||||
)
|
||||
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
|
||||
ray.get([trainer_init, inference_init])
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
|
||||
|
||||
dense_update_info, dense_payload_bytes = ray.get(
|
||||
train_model.get_dense_update_info.remote()
|
||||
)
|
||||
_, selected_token_ids, patch_digest, _ = ray.get(
|
||||
train_model.prepare_sparse_patch.remote(PROMPTS)
|
||||
)
|
||||
|
||||
inference_update = llm.update_weights.remote(
|
||||
dict(update_info=dense_update_info)
|
||||
)
|
||||
dense_send_ms, _ = ray.get(
|
||||
[
|
||||
train_model.broadcast_weights.remote(packed=False),
|
||||
inference_update,
|
||||
]
|
||||
)
|
||||
ray.get(llm.finish_weight_update.remote())
|
||||
ray.get(llm.wake_up.remote(tags=["scheduling"]))
|
||||
|
||||
dense_after = collect_vllm_generations(llm)
|
||||
|
||||
return {
|
||||
"dense_before": dense_before,
|
||||
"dense_after": dense_after,
|
||||
"selected_token_ids": selected_token_ids,
|
||||
"patch_digest": patch_digest,
|
||||
"dense_payload_bytes": dense_payload_bytes,
|
||||
"dense_send_ms": dense_send_ms,
|
||||
}
|
||||
finally:
|
||||
ray.kill(llm)
|
||||
|
||||
|
||||
def run_sparse_phase(
|
||||
train_model,
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
) -> dict[str, object]:
|
||||
ray.get(train_model.reset_model.remote())
|
||||
llm = launch_llm(scheduling_inference)
|
||||
try:
|
||||
sparse_before = collect_vllm_generations(llm)
|
||||
|
||||
ray.get(llm.sleep.remote(level=0))
|
||||
master_address, master_port = ray.get(train_model.create_rendezvous.remote())
|
||||
world_size = ray.get(llm.get_world_size.remote()) + 1
|
||||
inference_init = llm.init_weight_transfer_engine.remote(
|
||||
dict(
|
||||
init_info=dict(
|
||||
master_address=master_address,
|
||||
master_port=master_port,
|
||||
rank_offset=1,
|
||||
world_size=world_size,
|
||||
)
|
||||
)
|
||||
)
|
||||
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
|
||||
ray.get([trainer_init, inference_init])
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=False))
|
||||
|
||||
sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = (
|
||||
ray.get(train_model.prepare_sparse_patch.remote(PROMPTS))
|
||||
)
|
||||
|
||||
inference_update = llm.update_weights.remote(
|
||||
dict(update_info=sparse_update_info)
|
||||
)
|
||||
sparse_send_ms, _ = ray.get(
|
||||
[
|
||||
train_model.broadcast_pending_sparse_patch.remote(),
|
||||
inference_update,
|
||||
]
|
||||
)
|
||||
ray.get(llm.finish_weight_update.remote())
|
||||
ray.get(llm.wake_up.remote(tags=["scheduling"]))
|
||||
|
||||
sparse_after = collect_vllm_generations(llm)
|
||||
|
||||
return {
|
||||
"sparse_before": sparse_before,
|
||||
"sparse_after": sparse_after,
|
||||
"selected_token_ids": selected_token_ids,
|
||||
"patch_digest": patch_digest,
|
||||
"sparse_payload_bytes": sparse_payload_bytes,
|
||||
"sparse_send_ms": sparse_send_ms,
|
||||
}
|
||||
finally:
|
||||
ray.kill(llm)
|
||||
|
||||
|
||||
ray.init()
|
||||
|
||||
try:
|
||||
train_model = TrainModel.remote(MODEL_NAME)
|
||||
|
||||
pg_inference = placement_group([{"GPU": 1, "CPU": 0}])
|
||||
ray.get(pg_inference.ready())
|
||||
scheduling_inference = PlacementGroupSchedulingStrategy(
|
||||
placement_group=pg_inference,
|
||||
placement_group_capture_child_tasks=True,
|
||||
placement_group_bundle_index=0,
|
||||
)
|
||||
|
||||
dense_results = run_dense_phase(train_model, scheduling_inference)
|
||||
sparse_results = run_sparse_phase(train_model, scheduling_inference)
|
||||
|
||||
baseline_equal = token_sequences_match(
|
||||
dense_results["dense_before"],
|
||||
sparse_results["sparse_before"],
|
||||
)
|
||||
patch_selection_equal = (
|
||||
dense_results["selected_token_ids"] == sparse_results["selected_token_ids"]
|
||||
)
|
||||
patch_digest_equal = dense_results["patch_digest"] == sparse_results["patch_digest"]
|
||||
after_equal = token_sequences_match(
|
||||
dense_results["dense_after"],
|
||||
sparse_results["sparse_after"],
|
||||
)
|
||||
any_output_changed = any(
|
||||
before["token_ids"] != after["token_ids"]
|
||||
for before, after in zip(
|
||||
dense_results["dense_before"],
|
||||
dense_results["dense_after"],
|
||||
)
|
||||
)
|
||||
dense_payload_mb = dense_results["dense_payload_bytes"] / (1024 * 1024)
|
||||
sparse_payload_mb = sparse_results["sparse_payload_bytes"] / (1024 * 1024)
|
||||
|
||||
print_generations(
|
||||
"Dense baseline outputs",
|
||||
PROMPTS,
|
||||
dense_results["dense_before"],
|
||||
)
|
||||
print_generations(
|
||||
"Sparse baseline outputs", PROMPTS, sparse_results["sparse_before"]
|
||||
)
|
||||
print_generations(
|
||||
"Dense outputs after update", PROMPTS, dense_results["dense_after"]
|
||||
)
|
||||
print_generations(
|
||||
"Sparse outputs after update",
|
||||
PROMPTS,
|
||||
sparse_results["sparse_after"],
|
||||
)
|
||||
|
||||
print(f"patched_token_ids = {dense_results['selected_token_ids']}")
|
||||
print(f"patch_selection_equal = {patch_selection_equal}")
|
||||
print(f"dense_patch_digest = {dense_results['patch_digest']}")
|
||||
print(f"sparse_patch_digest = {sparse_results['patch_digest']}")
|
||||
print(f"patch_digest_equal = {patch_digest_equal}")
|
||||
print(f"baseline_equal = {baseline_equal}")
|
||||
print(f"after_equal = {after_equal}")
|
||||
print(f"any_output_changed = {any_output_changed}")
|
||||
print(f"dense_payload_mb = {dense_payload_mb:.2f}")
|
||||
print(f"sparse_payload_mb = {sparse_payload_mb:.2f}")
|
||||
print(f"dense_send_ms = {dense_results['dense_send_ms']:.2f}")
|
||||
print(f"sparse_send_ms = {sparse_results['sparse_send_ms']:.2f}")
|
||||
|
||||
if not baseline_equal:
|
||||
raise RuntimeError(
|
||||
"Dense and sparse phases did not start from the same baseline"
|
||||
)
|
||||
if not patch_selection_equal:
|
||||
raise RuntimeError("Dense and sparse phases used different sparse patches")
|
||||
if not patch_digest_equal:
|
||||
raise RuntimeError("Dense and sparse phases produced different patch values")
|
||||
if not after_equal:
|
||||
raise RuntimeError("Dense and sparse updates produced different outputs")
|
||||
if not any_output_changed:
|
||||
raise RuntimeError("Patch did not change the observed outputs")
|
||||
finally:
|
||||
ray.shutdown()
|
||||
@@ -233,7 +233,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, gemma4, glm45, glm47, hermes, hy_v3, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
|
||||
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -5,9 +5,9 @@ use std::sync::LazyLock;
|
||||
pub use vllm_tool_parser::{
|
||||
DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
|
||||
Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, HyV3ToolParser,
|
||||
KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MistralToolParser,
|
||||
Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError,
|
||||
ToolParserOutput,
|
||||
Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser,
|
||||
MistralToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser,
|
||||
ToolParserError, ToolParserOutput,
|
||||
};
|
||||
|
||||
use crate::parser::ParserFactory;
|
||||
@@ -24,6 +24,9 @@ pub mod names {
|
||||
pub const GEMMA4: &str = "gemma4";
|
||||
pub const HERMES: &str = "hermes";
|
||||
pub const HY_V3: &str = "hy_v3";
|
||||
// Matches the Python CLI name `--tool-call-parser internlm`, which Python
|
||||
// also routes to `Internlm2ToolParser` despite the version-agnostic name.
|
||||
pub const INTERNLM: &str = "internlm";
|
||||
pub const KIMI_K2: &str = "kimi_k2";
|
||||
pub const LLAMA3_JSON: &str = "llama3_json";
|
||||
pub const LLAMA4_JSON: &str = "llama4_json";
|
||||
@@ -62,6 +65,7 @@ impl ToolParserFactory {
|
||||
.register_parser::<Gemma4ToolParser>(names::GEMMA4)
|
||||
.register_parser::<HermesToolParser>(names::HERMES)
|
||||
.register_parser::<HyV3ToolParser>(names::HY_V3)
|
||||
.register_parser::<Internlm2ToolParser>(names::INTERNLM)
|
||||
.register_parser::<KimiK2ToolParser>(names::KIMI_K2)
|
||||
.register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
|
||||
.register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
|
||||
@@ -80,6 +84,12 @@ impl ToolParserFactory {
|
||||
.register_pattern("hermes", names::HERMES)
|
||||
.register_pattern("hy3", names::HY_V3)
|
||||
.register_pattern("hy_v3", names::HY_V3)
|
||||
// Narrow to `internlm2` substring so it matches `internlm2-chat-7b`
|
||||
// and `internlm2_5-7b-chat` but NOT `internlm-chat-7b` (InternLM v1,
|
||||
// routes to Llama), `internlm3-*` (also Llama-architecture per
|
||||
// vllm/model_executor/models/registry.py:146), or `Intern-S1` /
|
||||
// `Intern-S1-Pro` (separate intern-s1 parser, see PR #40115).
|
||||
.register_pattern("internlm2", names::INTERNLM)
|
||||
.register_pattern("llama-4", names::LLAMA4_JSON)
|
||||
.register_pattern("llama-3.2", names::LLAMA3_JSON)
|
||||
.register_pattern("llama-3.1", names::LLAMA3_JSON)
|
||||
|
||||
@@ -161,4 +161,33 @@ fn factory_new_resolves_default_patterns() {
|
||||
factory.resolve_name_for_model("org/mm-m2-base"),
|
||||
Some(names::MINIMAX_M2)
|
||||
);
|
||||
|
||||
// InternLM2 positive: both dashed and underscored versioned names route.
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("internlm/internlm2-chat-7b"),
|
||||
Some(names::INTERNLM)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("internlm/internlm2_5-7b-chat"),
|
||||
Some(names::INTERNLM)
|
||||
);
|
||||
|
||||
// Negative: other internlm-org models do NOT route to the InternLM2 parser,
|
||||
// since they use unrelated prompt formats.
|
||||
// - InternLM v1 (`internlm-chat-7b`) routes to Llama
|
||||
// - InternLM3 (`internlm3-8b-instruct`) routes to Llama
|
||||
// - Intern-S1 / Intern-S1-Pro have their own parser (Python PR #40115)
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("internlm/internlm-chat-7b"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("internlm/internlm3-8b-instruct"),
|
||||
None
|
||||
);
|
||||
assert_eq!(factory.resolve_name_for_model("internlm/Intern-S1"), None);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("internlm/Intern-S1-Pro"),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
@@ -20,6 +20,7 @@ use vllm_chat::{
|
||||
use vllm_text::{DecodedTextEvent, Finished, Prompt};
|
||||
|
||||
/// One model/parser configuration used to run the fixed roundtrip fixtures.
|
||||
#[derive(Clone)]
|
||||
struct RoundtripCase {
|
||||
/// Hugging Face model id resolved through the production backend loader.
|
||||
model_id: &'static str,
|
||||
@@ -31,11 +32,45 @@ struct RoundtripCase {
|
||||
tool_call_parser: ParserSelection,
|
||||
/// Reasoning parser selection used by the output processor.
|
||||
reasoning_parser: ParserSelection,
|
||||
/// How this model's chat template handles thinking mode.
|
||||
thinking_behavior: ThinkingBehavior,
|
||||
/// JSON formatting expected after this model's template has materialized
|
||||
/// tool-call arguments.
|
||||
json_fmt: JsonFmt,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum ThinkingBehavior {
|
||||
/// The chat template accepts explicit thinking on/off kwargs, and uses
|
||||
/// `default` when the request does not specify either kwarg.
|
||||
Toggleable { default: bool },
|
||||
/// The chat template always behaves as `value` for this fixture.
|
||||
Always { value: bool },
|
||||
}
|
||||
|
||||
impl ThinkingBehavior {
|
||||
fn default(self) -> bool {
|
||||
match self {
|
||||
Self::Toggleable { default } => default,
|
||||
Self::Always { value } => value,
|
||||
}
|
||||
}
|
||||
|
||||
fn fixtures(self) -> Vec<Option<bool>> {
|
||||
match self {
|
||||
Self::Toggleable { .. } => vec![
|
||||
Some(true), // explicitly enable thinking
|
||||
Some(false), // explicitly disable thinking
|
||||
None, // use default template behavior
|
||||
],
|
||||
Self::Always { value } => vec![
|
||||
Some(value), // explicitly request the supported thinking behavior
|
||||
None, // use default template behavior
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RoundtripCase {
|
||||
/// Qwen3 XML tool-call format with `qwen3` reasoning tags.
|
||||
fn qwen3() -> Self {
|
||||
@@ -44,6 +79,7 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "<|im_end|>\n",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
|
||||
json_fmt: spaced_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -55,6 +91,7 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "<|im_end|>\n",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -66,6 +103,7 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "[e~[\n",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Always { value: true },
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -77,6 +115,7 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "<|end▁of▁sentence|>",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Toggleable { default: false },
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -88,6 +127,7 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -100,6 +140,7 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "<|im_end|>",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
|
||||
json_fmt: spaced_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -135,35 +176,44 @@ roundtrip_tests! {
|
||||
|
||||
/// Run the fixed reasoning+content fixture for one model/parser case.
|
||||
async fn run_roundtrip_reasoning_and_content(case: RoundtripCase) -> Result<()> {
|
||||
for thinking in case.thinking_behavior.fixtures() {
|
||||
run_roundtrip_reasoning_and_content_inner(case.clone(), thinking).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_roundtrip_reasoning_and_content_inner(
|
||||
case: RoundtripCase,
|
||||
thinking: Option<bool>,
|
||||
) -> Result<()> {
|
||||
let backends = load_roundtrip_backends(&case).await?;
|
||||
let request = roundtrip_request(
|
||||
"roundtrip-reasoning-content",
|
||||
vec![ChatMessage::text(ChatRole::User, "What is 2 + 2?")],
|
||||
Vec::new(),
|
||||
thinking,
|
||||
);
|
||||
let expected_reasoning = "Need compute 2 + 2 directly.";
|
||||
let expected_text = "The answer is 4.";
|
||||
let effective_thinking = thinking.unwrap_or(case.thinking_behavior.default());
|
||||
|
||||
let result = run_roundtrip(
|
||||
&case,
|
||||
&backends,
|
||||
&request,
|
||||
AssistantMessage {
|
||||
content: vec![
|
||||
AssistantContentBlock::Reasoning {
|
||||
text: expected_reasoning.to_string(),
|
||||
},
|
||||
AssistantContentBlock::Text {
|
||||
text: expected_text.to_string(),
|
||||
},
|
||||
],
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
let assistant = {
|
||||
let mut content = Vec::new();
|
||||
if effective_thinking {
|
||||
content.push(AssistantContentBlock::Reasoning {
|
||||
text: expected_reasoning.to_string(),
|
||||
});
|
||||
}
|
||||
content.push(AssistantContentBlock::Text {
|
||||
text: expected_text.to_string(),
|
||||
});
|
||||
AssistantMessage { content }
|
||||
};
|
||||
let result = run_roundtrip(&case, &backends, &request, assistant).await?;
|
||||
|
||||
assert_eq!(
|
||||
result.parsed_message.reasoning().as_deref().map(str::trim),
|
||||
Some(expected_reasoning)
|
||||
effective_thinking.then_some(expected_reasoning)
|
||||
);
|
||||
assert_eq!(result.parsed_message.text().trim(), expected_text);
|
||||
assert_eq!(result.parsed_message.tool_calls().count(), 0);
|
||||
@@ -186,6 +236,7 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
|
||||
"Check Shanghai weather and add 1.00 plus 2.",
|
||||
)],
|
||||
test_tools(),
|
||||
Some(true), // always enable thinking in this fixture
|
||||
);
|
||||
let expected_reasoning = "Need call the weather and add tools.";
|
||||
let expected_text = "I will call the tools.";
|
||||
@@ -487,6 +538,7 @@ fn roundtrip_request(
|
||||
request_id: impl Into<String>,
|
||||
messages: Vec<ChatMessage>,
|
||||
tools: Vec<ChatTool>,
|
||||
thinking: Option<bool>,
|
||||
) -> ChatRequest {
|
||||
let mut request = ChatRequest {
|
||||
request_id: request_id.into(),
|
||||
@@ -500,10 +552,12 @@ fn roundtrip_request(
|
||||
..ChatRequest::for_test()
|
||||
};
|
||||
|
||||
// Enable thinking for some models so that rendering and parsing the reasoning block is
|
||||
// exercised in the roundtrip.
|
||||
for key in ["thinking", "enable_thinking"] {
|
||||
request.chat_options.template_kwargs.insert(key.to_string(), true.into());
|
||||
// Explicitly enable or disable thinking so that rendering and parsing the reasoning block is
|
||||
// exercised or skipped in the roundtrip. If unspecified, use the default template behavior.
|
||||
if let Some(thinking) = thinking {
|
||||
for key in ["thinking", "enable_thinking"] {
|
||||
request.chat_options.template_kwargs.insert(key.to_string(), thinking.into());
|
||||
}
|
||||
}
|
||||
|
||||
request
|
||||
|
||||
@@ -165,6 +165,15 @@ pub struct SharedRuntimeArgs {
|
||||
#[serde(default)]
|
||||
pub enable_log_requests: bool,
|
||||
|
||||
/// If specified, API server will add X-Request-Id header to responses.
|
||||
#[arg(
|
||||
long,
|
||||
default_missing_value = "true",
|
||||
num_args = 0..=1
|
||||
)]
|
||||
#[serde(default)]
|
||||
pub enable_request_id_headers: bool,
|
||||
|
||||
/// Disable periodic logging of engine statistics (throughput, queue depth,
|
||||
/// cache usage).
|
||||
#[arg(long)]
|
||||
@@ -238,6 +247,7 @@ impl SharedRuntimeArgs {
|
||||
default_chat_template_kwargs: self.default_chat_template_kwargs,
|
||||
chat_template_content_format: self.chat_template_content_format,
|
||||
enable_log_requests: self.enable_log_requests,
|
||||
enable_request_id_headers: self.enable_request_id_headers,
|
||||
disable_log_stats: self.disable_log_stats,
|
||||
grpc_port: self.grpc_port,
|
||||
shutdown_timeout,
|
||||
@@ -278,6 +288,7 @@ impl SharedRuntimeArgs {
|
||||
default_chat_template_kwargs: self.default_chat_template_kwargs,
|
||||
chat_template_content_format: self.chat_template_content_format,
|
||||
enable_log_requests: self.enable_log_requests,
|
||||
enable_request_id_headers: self.enable_request_id_headers,
|
||||
disable_log_stats: self.disable_log_stats,
|
||||
grpc_port: self.grpc_port,
|
||||
shutdown_timeout,
|
||||
|
||||
@@ -43,6 +43,7 @@ fn serve_args_forward_python_flags_with_separator() {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
served_model_name: [],
|
||||
},
|
||||
@@ -116,6 +117,46 @@ fn serve_args_accept_explicit_deepseek_v32_renderer() {
|
||||
assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serve_passes_enable_request_id_headers_into_config() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"serve",
|
||||
"Qwen/Qwen3-0.6B",
|
||||
"--enable-request-id-headers",
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Serve(args) = cli.command else {
|
||||
panic!("expected serve args");
|
||||
};
|
||||
let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
|
||||
assert!(config.enable_request_id_headers);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frontend_args_json_passes_enable_request_id_headers_into_config() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"frontend",
|
||||
"--listen-fd",
|
||||
"3",
|
||||
"--input-address",
|
||||
"ipc:///tmp/input.sock",
|
||||
"--output-address",
|
||||
"ipc:///tmp/output.sock",
|
||||
"--args-json",
|
||||
r#"{"model_tag":"Qwen/Qwen3-0.6B","enable_request_id_headers":true}"#,
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Frontend(args) = cli.command else {
|
||||
panic!("expected frontend args");
|
||||
};
|
||||
let config = args.into_config();
|
||||
assert!(config.enable_request_id_headers);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serve_args_reject_unknown_renderer_value() {
|
||||
let error = Cli::try_parse_from([
|
||||
@@ -218,6 +259,7 @@ fn frontend_args_accept_json() {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
served_model_name: [],
|
||||
},
|
||||
@@ -616,6 +658,7 @@ fn serve_args_accept_handshake_aliases() {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
served_model_name: [],
|
||||
},
|
||||
@@ -733,6 +776,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
@@ -795,6 +839,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
@@ -872,6 +917,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
|
||||
@@ -620,15 +620,6 @@ pub struct ServerUnsupportedArgs {
|
||||
#[arg(long)]
|
||||
pub middleware: Option<Unsupported>,
|
||||
|
||||
/// If specified, API server will add X-Request-Id header to responses.
|
||||
#[arg(
|
||||
long,
|
||||
visible_alias = "no-enable-request-id-headers",
|
||||
default_missing_value = "true",
|
||||
num_args = 0..=1
|
||||
)]
|
||||
pub enable_request_id_headers: Option<Unsupported>,
|
||||
|
||||
/// Disable FastAPI's OpenAPI schema, Swagger UI, and ReDoc endpoint.
|
||||
#[arg(
|
||||
long,
|
||||
|
||||
@@ -68,6 +68,7 @@ async fn main() -> Result<()> {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: ChatTemplateContentFormatOption::Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: Duration::ZERO,
|
||||
|
||||
@@ -61,6 +61,8 @@ pub struct Config {
|
||||
pub chat_template_content_format: ChatTemplateContentFormatOption,
|
||||
/// Log a summary line for each completed request.
|
||||
pub enable_log_requests: bool,
|
||||
/// When `true`, set `X-Request-Id` on every HTTP response.
|
||||
pub enable_request_id_headers: bool,
|
||||
/// When `true`, suppress periodic stats logging (throughput, queue depth,
|
||||
/// cache usage).
|
||||
pub disable_log_stats: bool,
|
||||
|
||||
@@ -85,7 +85,9 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
|
||||
};
|
||||
|
||||
Ok(Arc::new(
|
||||
AppState::new(served_model_names, chat).with_log_requests(config.enable_log_requests),
|
||||
AppState::new(served_model_names, chat)
|
||||
.with_log_requests(config.enable_log_requests)
|
||||
.with_request_id_headers(config.enable_request_id_headers),
|
||||
))
|
||||
}
|
||||
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
mod load;
|
||||
mod metrics;
|
||||
mod request_id;
|
||||
|
||||
pub use load::track_server_load;
|
||||
pub use metrics::track_http_metrics;
|
||||
pub use request_id::set_request_id_header;
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
use axum::extract::Request;
|
||||
use axum::http::HeaderValue;
|
||||
use axum::http::header::HeaderName;
|
||||
use axum::middleware::Next;
|
||||
use axum::response::Response;
|
||||
use uuid::Uuid;
|
||||
|
||||
const X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
|
||||
|
||||
/// Echo the request's `X-Request-Id` on the response, or generate a fresh
|
||||
/// `uuid4` hex if the request did not provide one.
|
||||
///
|
||||
/// Original Python:
|
||||
/// `vllm.entrypoints.openai.server_utils.XRequestIdMiddleware`.
|
||||
pub async fn set_request_id_header(req: Request, next: Next) -> Response {
|
||||
let incoming = req.headers().get(&X_REQUEST_ID).cloned();
|
||||
let mut response = next.run(req).await;
|
||||
let value = incoming.unwrap_or_else(|| {
|
||||
HeaderValue::from_str(&Uuid::new_v4().simple().to_string())
|
||||
.expect("uuid hex is valid header value")
|
||||
});
|
||||
response.headers_mut().insert(X_REQUEST_ID, value);
|
||||
response
|
||||
}
|
||||
@@ -56,11 +56,18 @@ fn build_router_with_dev_mode(state: Arc<AppState>, dev_mode_enabled: bool) -> R
|
||||
.route("/is_sleeping", get(sleep::is_sleeping))
|
||||
}
|
||||
|
||||
router
|
||||
let enable_request_id_headers = state.enable_request_id_headers;
|
||||
let mut router = router
|
||||
.with_state(state.clone())
|
||||
.layer(from_fn_with_state(state, middleware::track_server_load))
|
||||
.layer(from_fn(middleware::track_http_metrics))
|
||||
.layer(TraceLayer::new_for_http())
|
||||
.layer(TraceLayer::new_for_http());
|
||||
|
||||
if enable_request_id_headers {
|
||||
router = router.layer(from_fn(middleware::set_request_id_header));
|
||||
}
|
||||
|
||||
router
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -3,23 +3,33 @@ mod types;
|
||||
mod validate;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::convert::Infallible;
|
||||
use std::result::Result;
|
||||
use std::sync::Arc;
|
||||
|
||||
use asynk_strim_attr::{TryYielder, try_stream};
|
||||
use axum::Json;
|
||||
use axum::extract::State;
|
||||
use axum::http::HeaderMap;
|
||||
use axum::response::sse::{Event, Sse};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use futures::{Stream, StreamExt as _, pin_mut};
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tracing::info;
|
||||
use tracing::{error, info, trace};
|
||||
use tracing_futures::Instrument as _;
|
||||
use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
|
||||
use vllm_llm::{CollectedGenerateOutput, GenerateOutputStreamExt as _};
|
||||
use vllm_llm::{
|
||||
CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _,
|
||||
};
|
||||
|
||||
use self::convert::prepare_generate_request;
|
||||
use self::types::{GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice};
|
||||
use crate::error::{ApiError, server_error};
|
||||
use self::types::{
|
||||
GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice,
|
||||
GenerateResponseStreamChoice, GenerateStreamResponse,
|
||||
};
|
||||
use crate::error::{ApiError, bail_server_error, server_error};
|
||||
use crate::routes::openai::utils::logprobs::clamp_logprob;
|
||||
use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb};
|
||||
use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage};
|
||||
use crate::routes::openai::utils::validated_json::ValidatedJson;
|
||||
use crate::state::AppState;
|
||||
use crate::utils::resolve_request_context;
|
||||
@@ -46,6 +56,7 @@ pub async fn generate(
|
||||
let log_request = state.enable_log_requests;
|
||||
let include_logprobs = prepared.include_logprobs;
|
||||
let include_prompt_logprobs = prepared.include_prompt_logprobs;
|
||||
let stream = prepared.stream;
|
||||
|
||||
let raw_stream = match state
|
||||
.chat
|
||||
@@ -64,6 +75,20 @@ pub async fn generate(
|
||||
}
|
||||
};
|
||||
|
||||
if stream {
|
||||
let chunk_stream = generate_chunk_stream(
|
||||
raw_stream,
|
||||
prepared.request_id,
|
||||
log_request,
|
||||
prepared.include_usage,
|
||||
prepared.include_continuous_usage,
|
||||
include_logprobs,
|
||||
);
|
||||
let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
|
||||
|
||||
return Sse::new(sse_stream).into_response();
|
||||
}
|
||||
|
||||
let collected = match raw_stream.collect_output().instrument(request_span.clone()).await {
|
||||
Ok(collected) => collected,
|
||||
Err(error) => {
|
||||
@@ -98,6 +123,102 @@ pub async fn generate(
|
||||
Json(response).into_response()
|
||||
}
|
||||
|
||||
#[try_stream]
|
||||
async fn generate_chunk_stream(
|
||||
stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>>,
|
||||
request_id: String,
|
||||
log_request: bool,
|
||||
include_usage: bool,
|
||||
include_continuous_usage: bool,
|
||||
include_logprobs: bool,
|
||||
mut y: TryYielder<GenerateStreamResponse, ApiError>,
|
||||
) -> Result<(), ApiError> {
|
||||
pin_mut!(stream);
|
||||
let mut prompt_tokens: Option<u32> = None;
|
||||
let mut output_tokens = 0_u32;
|
||||
|
||||
while let Some(next) = stream.next().await {
|
||||
match next {
|
||||
Ok(output) => {
|
||||
if prompt_tokens.is_none() {
|
||||
prompt_tokens =
|
||||
output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len() as u32);
|
||||
}
|
||||
let usage_prompt_tokens = prompt_tokens.unwrap_or_default();
|
||||
|
||||
let token_ids = output.token_ids;
|
||||
output_tokens = output_tokens.saturating_add(token_ids.len() as u32);
|
||||
let finish_reason = output.finish_reason;
|
||||
|
||||
if matches!(finish_reason.as_ref(), Some(FinishReason::Error)) {
|
||||
bail_server_error!("Internal server error");
|
||||
}
|
||||
|
||||
if let Some(finish_reason) = finish_reason.as_ref()
|
||||
&& log_request
|
||||
{
|
||||
info!(
|
||||
stream = true,
|
||||
prompt_tokens = usage_prompt_tokens,
|
||||
output_tokens,
|
||||
finish_reason = finish_reason.as_str(),
|
||||
"generate finished"
|
||||
);
|
||||
}
|
||||
|
||||
if token_ids.is_empty() && finish_reason.is_none() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let logprobs = if include_logprobs && !token_ids.is_empty() {
|
||||
let logprobs = output.logprobs.as_ref().ok_or_else(|| {
|
||||
server_error!(
|
||||
"raw generate stream requested logprobs but generation returned none"
|
||||
)
|
||||
})?;
|
||||
Some(raw_logprobs_to_openai_chat(logprobs)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
y.yield_ok(GenerateStreamResponse {
|
||||
request_id: request_id.clone(),
|
||||
choices: vec![GenerateResponseStreamChoice {
|
||||
index: 0,
|
||||
logprobs,
|
||||
finish_reason: finish_reason.map(|reason| reason.as_str().to_string()),
|
||||
token_ids,
|
||||
}],
|
||||
usage: include_continuous_usage
|
||||
.then(|| Usage::from_counts(usage_prompt_tokens, output_tokens)),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
Err(error) => {
|
||||
error!(
|
||||
error = %error.as_report(),
|
||||
"raw generate stream failed"
|
||||
);
|
||||
bail_server_error!("{}", error.to_report_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if include_usage {
|
||||
y.yield_ok(GenerateStreamResponse {
|
||||
request_id,
|
||||
choices: Vec::new(),
|
||||
usage: Some(Usage::from_counts(
|
||||
prompt_tokens.unwrap_or_default(),
|
||||
output_tokens,
|
||||
)),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn collect_generate(
|
||||
collected: CollectedGenerateOutput,
|
||||
request_id: String,
|
||||
@@ -213,3 +334,94 @@ fn position_to_logprob_map(position: &PositionLogprobs) -> HashMap<u32, Generate
|
||||
fn format_token_id(token_id: u32) -> String {
|
||||
format!("token_id:{token_id}")
|
||||
}
|
||||
|
||||
/// Convert one raw-generate chunk stream into SSE events.
|
||||
#[try_stream]
|
||||
async fn generate_sse_stream(
|
||||
stream: impl Stream<Item = Result<GenerateStreamResponse, ApiError>>,
|
||||
mut y: TryYielder<Event, Infallible>,
|
||||
) -> Result<(), Infallible> {
|
||||
pin_mut!(stream);
|
||||
|
||||
while let Some(next) = stream.next().await {
|
||||
match next {
|
||||
Ok(chunk) => y.yield_ok(to_sse_event(&chunk)).await,
|
||||
Err(error) => {
|
||||
y.yield_ok(to_error_sse_event(&error)).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
y.yield_ok(done_sse_event()).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn to_sse_event(chunk: &GenerateStreamResponse) -> Event {
|
||||
let payload = serde_json::to_string(chunk).expect("generate chunk must serialize to JSON");
|
||||
trace!(payload, "generate emitting chunk");
|
||||
Event::default().data(payload)
|
||||
}
|
||||
|
||||
fn to_error_sse_event(error: &ApiError) -> Event {
|
||||
let payload = serde_json::to_string(&error.to_error_response())
|
||||
.expect("ErrorResponse must serialize to JSON");
|
||||
trace!(payload, "generate emitting error");
|
||||
Event::default().data(payload)
|
||||
}
|
||||
|
||||
fn done_sse_event() -> Event {
|
||||
trace!("generate emitting done");
|
||||
Event::default().data("[DONE]")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures::{TryStreamExt as _, stream};
|
||||
use vllm_llm::GeneratePromptInfo;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn generate_chunk_stream_captures_late_prompt_info() {
|
||||
let stream = stream::iter(vec![
|
||||
Ok(GenerateOutput {
|
||||
request_id: String::new(),
|
||||
prompt_info: None,
|
||||
token_ids: Vec::new(),
|
||||
logprobs: None,
|
||||
finish_reason: None,
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
Ok(GenerateOutput {
|
||||
request_id: String::new(),
|
||||
prompt_info: Some(GeneratePromptInfo {
|
||||
prompt_token_ids: Arc::from([11_u32, 22_u32]),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
token_ids: vec![33],
|
||||
logprobs: None,
|
||||
finish_reason: Some(FinishReason::stop_eos()),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let chunks: Vec<_> =
|
||||
generate_chunk_stream(stream, "raw-stream".to_string(), false, true, true, false)
|
||||
.try_collect()
|
||||
.await
|
||||
.expect("collect chunks");
|
||||
|
||||
assert_eq!(chunks.len(), 2);
|
||||
assert_eq!(
|
||||
chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens,
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
chunks[1].usage.as_ref().expect("final usage").prompt_tokens,
|
||||
2
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,6 +10,9 @@ use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
|
||||
pub struct PreparedRequest {
|
||||
pub request_id: String,
|
||||
pub text_request: TextRequest,
|
||||
pub stream: bool,
|
||||
pub include_usage: bool,
|
||||
pub include_continuous_usage: bool,
|
||||
pub include_logprobs: bool,
|
||||
pub include_prompt_logprobs: bool,
|
||||
}
|
||||
@@ -23,6 +26,18 @@ pub fn prepare_generate_request(
|
||||
) -> Result<PreparedRequest, ApiError> {
|
||||
validate::validate_request_compat(&request, served_model_names)?;
|
||||
|
||||
let stream = request.stream;
|
||||
let include_usage = request
|
||||
.stream_options
|
||||
.as_ref()
|
||||
.and_then(|options| options.include_usage)
|
||||
.unwrap_or(false);
|
||||
let include_continuous_usage = include_usage
|
||||
&& request
|
||||
.stream_options
|
||||
.as_ref()
|
||||
.and_then(|options| options.continuous_usage_stats)
|
||||
.unwrap_or(false);
|
||||
let include_logprobs = request.sampling_params.logprobs.is_some();
|
||||
let include_prompt_logprobs = request.sampling_params.prompt_logprobs.is_some();
|
||||
let mut sampling_params = request.sampling_params;
|
||||
@@ -47,6 +62,9 @@ pub fn prepare_generate_request(
|
||||
Ok(PreparedRequest {
|
||||
request_id: ctx.request_id,
|
||||
text_request,
|
||||
stream,
|
||||
include_usage,
|
||||
include_continuous_usage,
|
||||
include_logprobs,
|
||||
include_prompt_logprobs,
|
||||
})
|
||||
@@ -109,4 +127,28 @@ mod tests {
|
||||
Some(json!({"connector": "x"}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_generate_request_gates_continuous_usage_on_include_usage() {
|
||||
let request: GenerateRequest = serde_json::from_value(json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"token_ids": [11, 22],
|
||||
"stream": true,
|
||||
"stream_options": {
|
||||
"continuous_usage_stats": true
|
||||
},
|
||||
"sampling_params": {}
|
||||
}))
|
||||
.expect("parse request");
|
||||
|
||||
let prepared = prepare_generate_request(
|
||||
request,
|
||||
&["Qwen/Qwen1.5-0.5B-Chat".to_string()],
|
||||
ResolvedRequestContext::default(),
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
assert!(!prepared.include_usage);
|
||||
assert!(!prepared.include_continuous_usage);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,7 +5,7 @@ use serde_json::{Map, Value};
|
||||
use validator::Validate;
|
||||
use vllm_text::SamplingParams;
|
||||
|
||||
use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable};
|
||||
use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage};
|
||||
|
||||
/// vLLM-compatible request type for the token-in/token-out generate API.
|
||||
#[serde_with::skip_serializing_none]
|
||||
@@ -17,6 +17,7 @@ pub struct GenerateRequest {
|
||||
pub sampling_params: SamplingParams,
|
||||
#[serde(default)]
|
||||
pub stream: bool,
|
||||
pub stream_options: Option<StreamOptions>,
|
||||
pub cache_salt: Option<String>,
|
||||
#[serde(default)]
|
||||
pub priority: i32,
|
||||
@@ -37,6 +38,25 @@ pub(super) struct GenerateResponseChoice {
|
||||
pub token_ids: Vec<u32>,
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `GenerateResponseStreamChoice` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct GenerateResponseStreamChoice {
|
||||
pub index: u32,
|
||||
pub logprobs: Option<ChatLogProbs>,
|
||||
pub finish_reason: Option<String>,
|
||||
pub token_ids: Vec<u32>,
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `GenerateStreamResponse` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct GenerateStreamResponse {
|
||||
pub request_id: String,
|
||||
pub choices: Vec<GenerateResponseStreamChoice>,
|
||||
pub usage: Option<Usage>,
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `GenerateResponse` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
|
||||
@@ -13,8 +13,11 @@ pub(super) fn validate_request_compat(
|
||||
return Err(ApiError::model_not_found(model.clone()));
|
||||
}
|
||||
|
||||
if request.stream {
|
||||
bail_invalid_request!(param = "stream", "stream=true is not supported.");
|
||||
if request.stream_options.is_some() && !request.stream {
|
||||
bail_invalid_request!(
|
||||
param = "stream_options",
|
||||
"stream_options are only supported when stream=true."
|
||||
);
|
||||
}
|
||||
|
||||
if request.token_ids.is_empty() {
|
||||
@@ -65,11 +68,24 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_request_compat_rejects_streaming() {
|
||||
fn validate_request_compat_accepts_streaming() {
|
||||
let request = GenerateRequest {
|
||||
stream: true,
|
||||
..base_request()
|
||||
};
|
||||
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_request_compat_rejects_stream_options_without_streaming() {
|
||||
let request: GenerateRequest = serde_json::from_value(json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"token_ids": [11, 22],
|
||||
"stream": false,
|
||||
"stream_options": {"include_usage": true},
|
||||
"sampling_params": {}
|
||||
}))
|
||||
.expect("parse request");
|
||||
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
|
||||
}
|
||||
|
||||
|
||||
@@ -746,6 +746,20 @@ async fn test_app() -> axum::Router {
|
||||
)))
|
||||
}
|
||||
|
||||
async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {
|
||||
let (chat, engine_task) = test_models_with_engine_outputs_and_backend(
|
||||
b"engine-openai-request-id",
|
||||
default_stream_output_specs(),
|
||||
Arc::new(FakeChatBackend::new()),
|
||||
)
|
||||
.await;
|
||||
let app = build_router(Arc::new(
|
||||
AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat)
|
||||
.with_request_id_headers(true),
|
||||
));
|
||||
(app, engine_task)
|
||||
}
|
||||
|
||||
async fn test_health_app_with_engine_script<F>(
|
||||
script: F,
|
||||
) -> (axum::Router, Arc<AppState>, MockEngineTask)
|
||||
@@ -947,6 +961,18 @@ async fn health_status(app: &axum::Router) -> (StatusCode, Bytes) {
|
||||
(status, body)
|
||||
}
|
||||
|
||||
async fn health_response(app: &axum::Router, request_id: Option<&str>) -> axum::response::Response {
|
||||
let mut builder = Request::builder().method("GET").uri("/health");
|
||||
if let Some(request_id) = request_id {
|
||||
builder = builder.header("X-Request-Id", request_id);
|
||||
}
|
||||
|
||||
app.clone()
|
||||
.call(builder.body(Body::empty()).expect("build request"))
|
||||
.await
|
||||
.expect("call app")
|
||||
}
|
||||
|
||||
fn metric_value(rendered: &str, metric: &str, labels: Option<&str>) -> Option<f64> {
|
||||
rendered.lines().find_map(|line| {
|
||||
let rest = line.strip_prefix(metric)?;
|
||||
@@ -994,6 +1020,43 @@ async fn list_models_returns_configured_model() {
|
||||
assert_eq!(json["data"][0]["id"], "Qwen/Qwen1.5-0.5B-Chat");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn request_id_header_is_absent_by_default() {
|
||||
let app = test_app().await;
|
||||
let response = health_response(&app, None).await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert!(!response.headers().contains_key("x-request-id"));
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn request_id_header_generates_uuid_hex_when_enabled() {
|
||||
let (app, _engine_task) = test_app_with_request_id_headers().await;
|
||||
let response = health_response(&app, None).await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let request_id = response
|
||||
.headers()
|
||||
.get("x-request-id")
|
||||
.expect("x-request-id header")
|
||||
.to_str()
|
||||
.expect("header is ascii");
|
||||
assert_eq!(request_id.len(), 32);
|
||||
assert!(request_id.chars().all(|ch| ch.is_ascii_hexdigit()));
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn request_id_header_echoes_incoming_header_when_enabled() {
|
||||
let (app, _engine_task) = test_app_with_request_id_headers().await;
|
||||
let response = health_response(&app, Some("req-123")).await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(response.headers().get("x-request-id").unwrap(), "req-123");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn version_returns_engine_vllm_version() {
|
||||
@@ -2425,8 +2488,72 @@ async fn non_stream_raw_generate_returns_token_output_envelope() {
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn raw_generate_rejects_streaming() {
|
||||
let mut app = test_app().await;
|
||||
async fn stream_raw_generate_returns_sse_chunks_and_usage() {
|
||||
let ipc = IpcNamespace::new().expect("create ipc namespace");
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
let engine_id = b"engine-raw-generate-stream".to_vec();
|
||||
|
||||
let engine_task = MockEngineTask::new(spawn_mock_engine_task(
|
||||
handshake_address.clone(),
|
||||
engine_id.clone(),
|
||||
|dealer, push| {
|
||||
boxed_test_future(async move {
|
||||
let add = recv_engine_message(dealer).await;
|
||||
let request: EngineCoreRequest =
|
||||
rmp_serde::from_slice(&add[1]).expect("decode request");
|
||||
assert_eq!(request.prompt_token_ids.as_deref(), Some(&[11, 22][..]));
|
||||
assert_eq!(request.external_req_id.as_deref(), Some("raw-stream"));
|
||||
|
||||
send_outputs(
|
||||
push,
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![
|
||||
request_output_with_logprobs(
|
||||
&request.request_id,
|
||||
vec![33],
|
||||
None,
|
||||
None,
|
||||
Some(sample_logprobs_for_token(33, 34)),
|
||||
None,
|
||||
),
|
||||
request_output_with_logprobs(
|
||||
&request.request_id,
|
||||
vec![44],
|
||||
Some(EngineCoreFinishReason::Stop),
|
||||
None,
|
||||
Some(sample_logprobs_for_token(44, 45)),
|
||||
None,
|
||||
),
|
||||
],
|
||||
scheduler_stats: None,
|
||||
timestamp: 0.0,
|
||||
utility_output: None,
|
||||
finished_requests: None,
|
||||
wave_complete: None,
|
||||
start_wave: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
},
|
||||
));
|
||||
|
||||
let client = EngineCoreClient::connect(
|
||||
EngineCoreClientConfig::new_single(handshake_address)
|
||||
.with_model_name("test-model")
|
||||
.with_local_input_output_addresses(
|
||||
Some(ipc.input_endpoint()),
|
||||
Some(ipc.output_endpoint()),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("connect client");
|
||||
let chat = ChatLlm::from_shared_backend(Llm::new(client), Arc::new(FakeChatBackend::new()));
|
||||
let mut app = build_router(Arc::new(AppState::new(
|
||||
vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
|
||||
chat,
|
||||
)));
|
||||
|
||||
let response = app
|
||||
.call(
|
||||
@@ -2437,9 +2564,17 @@ async fn raw_generate_rejects_streaming() {
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"request_id": "raw-stream",
|
||||
"token_ids": [11, 22],
|
||||
"stream": true,
|
||||
"sampling_params": {}
|
||||
"stream_options": {
|
||||
"include_usage": true,
|
||||
"continuous_usage_stats": true
|
||||
},
|
||||
"sampling_params": {
|
||||
"max_tokens": 2,
|
||||
"logprobs": 1
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
@@ -2448,10 +2583,196 @@ async fn raw_generate_rejects_streaming() {
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.headers().get("content-type").and_then(|value| value.to_str().ok()),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
|
||||
assert_eq!(json["error"]["param"], "stream");
|
||||
engine_task.await.expect("mock engine task");
|
||||
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
|
||||
let payloads = sse_data_payloads(&text);
|
||||
assert_eq!(payloads.len(), 4, "{text}");
|
||||
|
||||
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
|
||||
assert_eq!(first["request_id"], "raw-stream");
|
||||
assert_eq!(first["choices"][0]["index"], 0);
|
||||
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
|
||||
assert_eq!(
|
||||
first["choices"][0]["logprobs"]["content"][0]["token"],
|
||||
"token_id:33"
|
||||
);
|
||||
assert_eq!(first["usage"]["prompt_tokens"], 2);
|
||||
assert_eq!(first["usage"]["completion_tokens"], 1);
|
||||
|
||||
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
|
||||
assert_eq!(second["choices"][0]["token_ids"], json!([44]));
|
||||
assert_eq!(second["choices"][0]["finish_reason"], "stop");
|
||||
assert_eq!(second["usage"]["completion_tokens"], 2);
|
||||
|
||||
let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
|
||||
assert_eq!(usage["choices"], json!([]));
|
||||
assert_eq!(usage["usage"]["prompt_tokens"], 2);
|
||||
assert_eq!(usage["usage"]["completion_tokens"], 2);
|
||||
assert_eq!(usage["usage"]["total_tokens"], 4);
|
||||
assert_eq!(payloads[3], "[DONE]");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn stream_raw_generate_emits_final_usage_without_continuous_usage() {
|
||||
let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
|
||||
(vec![33], None),
|
||||
(vec![44], Some(EngineCoreFinishReason::Stop)),
|
||||
])
|
||||
.await;
|
||||
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/inference/v1/generate")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"request_id": "raw-stream-final-usage",
|
||||
"token_ids": [11, 22],
|
||||
"stream": true,
|
||||
"stream_options": {
|
||||
"include_usage": true
|
||||
},
|
||||
"sampling_params": {
|
||||
"max_tokens": 2
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
engine_task.await.expect("mock engine task");
|
||||
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
|
||||
let payloads = sse_data_payloads(&text);
|
||||
assert_eq!(payloads.len(), 4, "{text}");
|
||||
|
||||
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
|
||||
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
|
||||
assert!(first.get("usage").is_none());
|
||||
|
||||
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
|
||||
assert_eq!(second["choices"][0]["token_ids"], json!([44]));
|
||||
assert_eq!(second["choices"][0]["finish_reason"], "stop");
|
||||
assert!(second.get("usage").is_none());
|
||||
|
||||
let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
|
||||
assert_eq!(usage["choices"], json!([]));
|
||||
assert_eq!(usage["usage"]["prompt_tokens"], 2);
|
||||
assert_eq!(usage["usage"]["completion_tokens"], 2);
|
||||
assert_eq!(usage["usage"]["total_tokens"], 4);
|
||||
assert_eq!(payloads[3], "[DONE]");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn stream_raw_generate_emits_empty_finish_chunk() {
|
||||
let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
|
||||
(vec![33], None),
|
||||
(vec![], Some(EngineCoreFinishReason::Stop)),
|
||||
])
|
||||
.await;
|
||||
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/inference/v1/generate")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"request_id": "raw-stream-empty-finish",
|
||||
"token_ids": [11, 22],
|
||||
"stream": true,
|
||||
"sampling_params": {
|
||||
"max_tokens": 2
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
engine_task.await.expect("mock engine task");
|
||||
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
|
||||
let payloads = sse_data_payloads(&text);
|
||||
assert_eq!(payloads.len(), 3, "{text}");
|
||||
|
||||
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
|
||||
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
|
||||
assert!(first["choices"][0].get("finish_reason").is_none());
|
||||
|
||||
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
|
||||
assert_eq!(second["choices"][0]["token_ids"], json!([]));
|
||||
assert_eq!(second["choices"][0]["finish_reason"], "stop");
|
||||
assert_eq!(payloads[2], "[DONE]");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn stream_raw_generate_error_finish_returns_sse_error() {
|
||||
let (mut app, engine_task) =
|
||||
test_app_with_stream_output_specs(vec![(vec![], Some(EngineCoreFinishReason::Error))])
|
||||
.await;
|
||||
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/inference/v1/generate")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"request_id": "raw-stream-error",
|
||||
"token_ids": [11, 22],
|
||||
"stream": true,
|
||||
"stream_options": {
|
||||
"include_usage": true
|
||||
},
|
||||
"sampling_params": {
|
||||
"max_tokens": 2
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
engine_task.await.expect("mock engine task");
|
||||
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
|
||||
|
||||
assert!(text.contains("\"type\":\"server_error\""), "{text}");
|
||||
assert!(text.contains("Internal server error"), "{text}");
|
||||
assert!(!text.contains("\"finish_reason\":\"error\""), "{text}");
|
||||
assert!(!text.contains("\"usage\":"), "{text}");
|
||||
assert!(text.trim_end().ends_with("data: [DONE]"), "{text}");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
|
||||
@@ -17,6 +17,8 @@ pub struct AppState {
|
||||
pub chat: ChatLlm,
|
||||
/// Whether to log a summary line for each completed request.
|
||||
pub enable_log_requests: bool,
|
||||
/// Whether to set X-Request-Id on every HTTP response.
|
||||
pub enable_request_id_headers: bool,
|
||||
/// Number of in-flight inference requests currently owned by this frontend.
|
||||
server_load: AtomicU64,
|
||||
}
|
||||
@@ -39,6 +41,7 @@ impl AppState {
|
||||
served_model_names,
|
||||
chat,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
server_load: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
@@ -49,6 +52,12 @@ impl AppState {
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable X-Request-Id response headers.
|
||||
pub fn with_request_id_headers(mut self, enabled: bool) -> Self {
|
||||
self.enable_request_id_headers = enabled;
|
||||
self
|
||||
}
|
||||
|
||||
/// The primary model name echoed back in API responses (the first served
|
||||
/// name).
|
||||
pub fn primary_model_name(&self) -> &str {
|
||||
|
||||
@@ -104,7 +104,7 @@ impl DeepSeekDsmlToolParser {
|
||||
self.tool_parameters.convert_param_with_schema(
|
||||
&name,
|
||||
¶m.name,
|
||||
¶m.value,
|
||||
param.value,
|
||||
)
|
||||
};
|
||||
arguments.insert(param.name, value);
|
||||
|
||||
@@ -8,7 +8,7 @@ const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
marker_whitespace: JsonToolCallWhitespace::Optional,
|
||||
delimiter: None,
|
||||
name_key: "name",
|
||||
arguments_key: "arguments",
|
||||
arguments_key: &["arguments"],
|
||||
};
|
||||
|
||||
/// Tool parser for Hermes XML-wrapped JSON tool calls.
|
||||
|
||||
@@ -0,0 +1,351 @@
|
||||
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
|
||||
use crate::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
const INTERNLM2_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
parser_name: "InternLM2",
|
||||
start_marker: "<|action_start|><|plugin|>",
|
||||
end_marker: "<|action_end|>",
|
||||
marker_whitespace: JsonToolCallWhitespace::Optional,
|
||||
delimiter: None,
|
||||
name_key: "name",
|
||||
// The Python parser's `get_arguments()` accepts either `parameters` or
|
||||
// `arguments` and prefers `parameters` when both are present. This Rust
|
||||
// parser uses first-encountered semantics because the header parser only
|
||||
// permits one args key per tool-call object; if a future model emits
|
||||
// both keys in the same object, the Rust port will accept the first one
|
||||
// and reject the trailing one as a syntax error rather than silently
|
||||
// shadowing it.
|
||||
arguments_key: &["parameters", "arguments"],
|
||||
};
|
||||
|
||||
/// Tool parser for InternLM2 special-token wrapped JSON tool calls.
|
||||
///
|
||||
/// Example tool call content:
|
||||
///
|
||||
/// ```text
|
||||
/// <|action_start|><|plugin|>{"name": "get_weather", "parameters": {"location":"Tokyo"}}<|action_end|>
|
||||
/// ```
|
||||
///
|
||||
/// Arguments are already OpenAI-style JSON text, so they are streamed as raw
|
||||
/// argument deltas without schema conversion or JSON normalization.
|
||||
///
|
||||
/// # Divergences from the Python reference
|
||||
///
|
||||
/// This Rust port intentionally diverges from
|
||||
/// `vllm/tool_parsers/internlm2_tool_parser.py` in two user-visible ways:
|
||||
///
|
||||
/// - **Parallel tool calls are supported.** Python silently drops every
|
||||
/// `<|action_start|>` block after the first (`current_tool_id > 0` returns
|
||||
/// an empty delta); this parser emits every well-formed block with
|
||||
/// incrementing `tool_index`. Models that legitimately emit multiple action
|
||||
/// blocks therefore produce more tool calls under Rust than under Python.
|
||||
/// - **End-marker bytes inside JSON string values are preserved.** Python
|
||||
/// does `action.split("<|action_end|>")[0]` which truncates regardless of
|
||||
/// JSON context; this parser scans matched braces and quotes so a literal
|
||||
/// `<|action_end|>` inside an arguments string is forwarded intact.
|
||||
/// - **Only whitespace is allowed before the `{`.** Python's non-streaming
|
||||
/// `action[action.find("{"):]` drops any bytes before the first `{`, but
|
||||
/// its streaming path has no equivalent and the model format always emits
|
||||
/// `<|plugin|>{...`; this parser allows only whitespace there, matching the
|
||||
/// other JSON parsers in this crate.
|
||||
/// - **Truncated tool calls error rather than silently dropping.** Python's
|
||||
/// streaming wrapper swallows mid-stream errors with `except Exception:
|
||||
/// return None` (logging a traceback) while its non-streaming path raises
|
||||
/// `JSONDecodeError`; this parser returns an `incomplete InternLM2 tool
|
||||
/// call` error from `finish()`, matching the other JSON parsers and Python's
|
||||
/// non-streaming behavior.
|
||||
///
|
||||
/// # Known unaddressed divergences (TODO)
|
||||
///
|
||||
/// The following Python behaviors are NOT yet matched. They are deferred to
|
||||
/// follow-up work because they require non-local changes to the shared
|
||||
/// `JsonToolCallParser` core that would affect Hermes / Llama / Mistral /
|
||||
/// Qwen as well. If a real-world InternLM2 deployment hits one of these,
|
||||
/// prioritize the corresponding fix.
|
||||
///
|
||||
/// - **Arguments value type.** The shared core requires the arguments value
|
||||
/// to be a JSON object (`take_json_object` rejects anything not starting
|
||||
/// with `{`). Python's `json.dumps(action_dict.get("parameters", ...))`
|
||||
/// accepts `null`, arrays, strings, and numbers and round-trips them
|
||||
/// verbatim. Models that legitimately emit `"parameters":null` will hard-
|
||||
/// fail under Rust.
|
||||
/// - **Unknown arguments key.** Python falls back to `{}` via
|
||||
/// `action_dict.get("parameters", action_dict.get("arguments", {}))` when
|
||||
/// neither key is present; the Rust header parser raises
|
||||
/// `parsing failed: invalid InternLM2` for any unrecognized key. A model
|
||||
/// that emits a typo (e.g. `"params"`) breaks the whole response.
|
||||
/// - **Field order independence.** The header parser requires the JSON keys
|
||||
/// to appear in the order `name` then arguments key. Python's
|
||||
/// `json.loads` + `dict.get` is order-independent, so a model emitting
|
||||
/// `{"parameters":{...},"name":"foo"}` parses in Python but fails in Rust.
|
||||
pub struct Internlm2ToolParser {
|
||||
inner: JsonToolCallParser,
|
||||
}
|
||||
|
||||
impl Internlm2ToolParser {
|
||||
/// Create an InternLM2 tool parser.
|
||||
fn new(_tools: &[Tool]) -> Self {
|
||||
Self {
|
||||
inner: JsonToolCallParser::new(INTERNLM2_CONFIG),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolParser for Internlm2ToolParser {
|
||||
/// Create a boxed InternLM2 tool parser.
|
||||
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
/// Preserve special-token markers while decoding, since
|
||||
/// `<|action_start|>`, `<|plugin|>`, and `<|action_end|>` are tokenizer
|
||||
/// special tokens in InternLM2 models.
|
||||
fn preserve_special_tokens(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Feed one decoded text chunk through the InternLM2 parser.
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
self.inner.parse_into(chunk, output)
|
||||
}
|
||||
|
||||
/// Flush any buffered partial state at end of stream.
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
self.inner.finish()
|
||||
}
|
||||
|
||||
/// Clear parser state and return currently uncommitted buffered text.
|
||||
fn reset(&mut self) -> String {
|
||||
self.inner.reset()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use expect_test::expect;
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::Internlm2ToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
const ACTION_START: &str = "<|action_start|><|plugin|>";
|
||||
const ACTION_END: &str = "<|action_end|>";
|
||||
|
||||
fn build_tool_call(function_name: &str, args_key: &str, arguments: &str) -> String {
|
||||
format!(
|
||||
r#"{ACTION_START}{{"name":"{function_name}","{args_key}":{arguments}}}{ACTION_END}"#
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_parse_complete_without_tool_call_keeps_text() {
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
let result = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(result.normal_text, "Hello, world!");
|
||||
assert!(result.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_parse_complete_extracts_parameters_key() {
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
let arguments = r#"{"location":"Tokyo","days":"3"}"#;
|
||||
let result = parser
|
||||
.parse_complete(&format!(
|
||||
"Let me check.\n{}",
|
||||
build_tool_call("get_weather", "parameters", arguments)
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.normal_text, "Let me check.\n");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].tool_index, 0);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_parse_complete_extracts_arguments_key_fallback() {
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
let arguments = r#"{"location":"Tokyo"}"#;
|
||||
let result = parser
|
||||
.parse_complete(&build_tool_call("get_weather", "arguments", arguments))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_accepts_whitespace_after_plugin_marker() {
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
let result = parser
|
||||
.parse_complete(&format!(
|
||||
r#"{ACTION_START}
|
||||
{{"name":"get_weather","parameters":{{}}}}{ACTION_END}"#
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_does_not_validate_or_normalize_arguments() {
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
let arguments = r#"{"location":"Tokyo",}"#;
|
||||
let result = parser
|
||||
.parse_complete(&build_tool_call("get_weather", "parameters", arguments))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_streaming_emits_argument_deltas() {
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
let chunks = [
|
||||
"preface <|action",
|
||||
"_start|><|plugin|>",
|
||||
r#"{"name":"get_weather","parameters":"#,
|
||||
r#"{"location":"#,
|
||||
r#""Beijing""#,
|
||||
r#"}"#,
|
||||
r#"}<|action_end|> suffix"#,
|
||||
];
|
||||
|
||||
let mut result = ToolParserOutput::default();
|
||||
let mut observed_arguments = Vec::new();
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
);
|
||||
result.append(next);
|
||||
}
|
||||
result.append(parser.finish().unwrap());
|
||||
|
||||
assert_eq!(
|
||||
observed_arguments,
|
||||
[r#"{"location":"#, r#""Beijing""#, r#"}"#]
|
||||
);
|
||||
assert_eq!(result.normal_text, "preface suffix");
|
||||
assert_eq!(
|
||||
result.coalesce_calls().calls[0].arguments,
|
||||
r#"{"location":"Beijing"}"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_streaming_handles_split_markers() {
|
||||
let input = format!(
|
||||
"hello {}",
|
||||
build_tool_call("get_weather", "parameters", r#"{"location":"Tokyo"}"#)
|
||||
);
|
||||
let chunks = split_by_chars(&input, 5);
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(result.normal_text, "hello ");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_streaming_extracts_multiple_blocks() {
|
||||
let input = format!(
|
||||
"{}{}",
|
||||
build_tool_call("get_weather", "parameters", r#"{"location":"Shanghai"}"#),
|
||||
build_tool_call("add", "arguments", r#"{"x":1,"y":2}"#),
|
||||
);
|
||||
let chunks = split_by_chars(&input, 7);
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"location\":\"Shanghai\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"add",
|
||||
),
|
||||
arguments: "{\"x\":1,\"y\":2}",
|
||||
},
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
.assert_debug_eq(&result);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_keeps_end_marker_literal_inside_json_string() {
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
let arguments = format!(r#"{{"text":"literal {ACTION_END} inside"}}"#);
|
||||
let input = build_tool_call("echo", "parameters", &arguments);
|
||||
|
||||
let result = parser.parse_complete(&input).unwrap();
|
||||
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_finish_errors_on_truncated_tool_call() {
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
let pre_finish = parser
|
||||
.parse_chunk(&format!(
|
||||
r#"{ACTION_START}{{"name":"get_weather","parameters":{{"location""#
|
||||
))
|
||||
.unwrap();
|
||||
let error = parser.finish().unwrap_err();
|
||||
|
||||
assert_eq!(
|
||||
pre_finish.calls[0].name.as_deref(),
|
||||
Some("get_weather"),
|
||||
"name delta is still emitted from parse_chunk() before truncation",
|
||||
);
|
||||
assert!(
|
||||
error.to_report_string().contains("incomplete InternLM2 tool call"),
|
||||
"finish() reports the truncated tool call as incomplete: {}",
|
||||
error.to_report_string(),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_unknown_arguments_key_fails() {
|
||||
let mut parser = Internlm2ToolParser::new(&test_tools());
|
||||
let input = build_tool_call("get_weather", "params", r#"{"location":"Tokyo"}"#);
|
||||
|
||||
let error = parser.parse_chunk(&input).unwrap_err();
|
||||
|
||||
expect![[r#"
|
||||
tool parser parsing failed: invalid InternLM2
|
||||
expected `parameters`, `arguments`"#]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internlm2_preserve_special_tokens_is_true() {
|
||||
let parser = Internlm2ToolParser::new(&test_tools());
|
||||
assert!(parser.preserve_special_tokens());
|
||||
}
|
||||
}
|
||||
@@ -203,7 +203,7 @@ fn llama_tool_call_header_event(input: &mut JsonToolInput<'_>) -> ModalResult<Ll
|
||||
marker_whitespace: JsonToolCallWhitespace::Optional,
|
||||
delimiter: Some(";"),
|
||||
name_key: "name",
|
||||
arguments_key: "parameters",
|
||||
arguments_key: &["parameters"],
|
||||
};
|
||||
|
||||
match tool_call_header_event(input, CONFIG)? {
|
||||
|
||||
@@ -8,7 +8,7 @@ const MISTRAL_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
marker_whitespace: JsonToolCallWhitespace::Optional,
|
||||
delimiter: Some(","),
|
||||
name_key: "name",
|
||||
arguments_key: "arguments",
|
||||
arguments_key: &["arguments"],
|
||||
};
|
||||
|
||||
/// Tool parser for Mistral JSON-array tool calls.
|
||||
|
||||
@@ -1,20 +1,22 @@
|
||||
//! Shared parser core for JSON tool calls wrapped by text markers.
|
||||
|
||||
pub use hermes::HermesToolParser;
|
||||
pub use internlm2::Internlm2ToolParser;
|
||||
pub use llama::Llama3JsonToolParser;
|
||||
pub use mistral::MistralToolParser;
|
||||
pub use qwen::Qwen3XmlToolParser;
|
||||
|
||||
mod hermes;
|
||||
mod internlm2;
|
||||
mod llama;
|
||||
mod mistral;
|
||||
mod qwen;
|
||||
|
||||
use winnow::ascii::multispace0 as ws0;
|
||||
use winnow::combinator::{alt, seq};
|
||||
use winnow::error::{ModalResult, StrContext, StrContextValue};
|
||||
use winnow::error::{AddContext, ModalResult, StrContext, StrContextValue};
|
||||
use winnow::prelude::*;
|
||||
use winnow::stream::Partial;
|
||||
use winnow::stream::{Partial, Stream};
|
||||
use winnow::token::literal;
|
||||
|
||||
use super::utils::{
|
||||
@@ -32,7 +34,10 @@ struct JsonToolCallConfig {
|
||||
marker_whitespace: JsonToolCallWhitespace,
|
||||
delimiter: Option<&'static str>,
|
||||
name_key: &'static str,
|
||||
arguments_key: &'static str,
|
||||
/// Candidate JSON keys naming the arguments payload, tried in order.
|
||||
/// Most parsers use a single key like `["arguments"]`, but some accept
|
||||
/// multiple (e.g. InternLM2 accepts `parameters` or `arguments`).
|
||||
arguments_key: &'static [&'static str],
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
@@ -224,7 +229,7 @@ fn tool_call_header_event(
|
||||
_: ws0,
|
||||
_: literal(","),
|
||||
_: ws0,
|
||||
_: |input: &mut JsonToolInput<'_>| json_key(input, config.arguments_key),
|
||||
_: |input: &mut JsonToolInput<'_>| json_arguments_key(input, config.arguments_key),
|
||||
_: ws0,
|
||||
_: literal(":"),
|
||||
_: ws0,
|
||||
@@ -246,6 +251,39 @@ fn json_key(input: &mut JsonToolInput<'_>, key: &'static str) -> ModalResult<()>
|
||||
.parse_next(input)
|
||||
}
|
||||
|
||||
/// Parse a JSON object key accepting any of `candidates`.
|
||||
///
|
||||
/// The full quoted key is consumed and compared against the candidate list,
|
||||
/// so this works correctly under partial input regardless of key lengths.
|
||||
///
|
||||
/// On mismatch, each candidate is attached as its own `Expected` context so the
|
||||
/// error enumerates every valid key ("expected `a`, expected `b`"). Because
|
||||
/// `StrContextValue::StringLiteral` carries a single `&'static str`, the
|
||||
/// contexts are added in a loop over `candidates` rather than through chained
|
||||
/// `.context(...)` calls, which keeps the diagnostics complete for any number
|
||||
/// of candidates.
|
||||
fn json_arguments_key(
|
||||
input: &mut JsonToolInput<'_>,
|
||||
candidates: &'static [&'static str],
|
||||
) -> ModalResult<()> {
|
||||
let start = input.checkpoint();
|
||||
json_str
|
||||
.verify(|key: &String| candidates.contains(&key.as_str()))
|
||||
.void()
|
||||
.parse_next(input)
|
||||
.map_err(|err| {
|
||||
err.map(|context_error| {
|
||||
candidates.iter().fold(context_error, |context_error, candidate| {
|
||||
context_error.add_context(
|
||||
&*input,
|
||||
&start,
|
||||
StrContext::Expected(StrContextValue::StringLiteral(candidate)),
|
||||
)
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// Parse one event inside a marker-wrapped JSON tool-call arguments payload.
|
||||
fn parse_arguments_event(
|
||||
input: &mut JsonToolInput<'_>,
|
||||
@@ -341,7 +379,7 @@ mod tests {
|
||||
marker_whitespace: JsonToolCallWhitespace::Optional,
|
||||
delimiter: Some("<"),
|
||||
name_key: "function",
|
||||
arguments_key: "parameters",
|
||||
arguments_key: &["parameters"],
|
||||
};
|
||||
|
||||
fn build_tool_call(function_name: &str, arguments: &str) -> String {
|
||||
|
||||
@@ -8,7 +8,7 @@ const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
marker_whitespace: JsonToolCallWhitespace::Exact("\n"),
|
||||
delimiter: None,
|
||||
name_key: "name",
|
||||
arguments_key: "arguments",
|
||||
arguments_key: &["arguments"],
|
||||
};
|
||||
|
||||
/// Tool parser for Qwen XML-wrapped JSON tool calls.
|
||||
|
||||
@@ -24,7 +24,10 @@ pub use error::{Result, ToolParserError};
|
||||
pub use gemma4::Gemma4ToolParser;
|
||||
pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser};
|
||||
pub use hy_v3::HyV3ToolParser;
|
||||
pub use json::{HermesToolParser, Llama3JsonToolParser, MistralToolParser, Qwen3XmlToolParser};
|
||||
pub use json::{
|
||||
HermesToolParser, Internlm2ToolParser, Llama3JsonToolParser, MistralToolParser,
|
||||
Qwen3XmlToolParser,
|
||||
};
|
||||
pub use kimi_k2::KimiK2ToolParser;
|
||||
pub use minimax_m2::MinimaxM2ToolParser;
|
||||
pub use qwen_coder::Qwen3CoderToolParser;
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{Number, Value};
|
||||
use serde_json::{Map, Number, Value};
|
||||
|
||||
use crate::Tool;
|
||||
|
||||
@@ -21,6 +21,29 @@ pub(super) struct ToolSchema {
|
||||
params: BTreeMap<String, JsonParamType>,
|
||||
}
|
||||
|
||||
/// Parameter input for schema-aware conversion.
|
||||
///
|
||||
/// It can be either a raw text string, or a structured input with named child elements.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(super) enum ParamInput {
|
||||
Text(String),
|
||||
#[allow(dead_code)]
|
||||
Elements(Vec<ParamElement>),
|
||||
}
|
||||
|
||||
impl From<String> for ParamInput {
|
||||
fn from(value: String) -> Self {
|
||||
Self::Text(value)
|
||||
}
|
||||
}
|
||||
|
||||
/// One named structured parameter child.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(super) struct ParamElement {
|
||||
pub name: String,
|
||||
pub value: ParamInput,
|
||||
}
|
||||
|
||||
/// Normalized JSON parameter type used for raw string coercion.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(super) enum JsonParamType {
|
||||
@@ -28,8 +51,13 @@ pub(super) enum JsonParamType {
|
||||
Integer,
|
||||
Number,
|
||||
Boolean,
|
||||
Object,
|
||||
Array,
|
||||
Object {
|
||||
properties: BTreeMap<String, JsonParamType>,
|
||||
additional_properties: Option<Box<JsonParamType>>,
|
||||
},
|
||||
Array {
|
||||
items: Option<Box<JsonParamType>>,
|
||||
},
|
||||
Null,
|
||||
OneOf(Vec<JsonParamType>),
|
||||
}
|
||||
@@ -45,33 +73,39 @@ impl ToolSchemas {
|
||||
Self { tools }
|
||||
}
|
||||
|
||||
/// Convert raw string parameter values for one named tool.
|
||||
/// Convert parameter values for one named tool.
|
||||
///
|
||||
/// Unknown tool names use an empty schema, so all parameters fall back to
|
||||
/// strings.
|
||||
pub(super) fn convert_params_with_schema(
|
||||
/// strings or object-like JSON for structured inputs.
|
||||
pub(super) fn convert_params_with_schema<P>(
|
||||
&self,
|
||||
function_name: &str,
|
||||
params: Vec<(String, String)>,
|
||||
) -> serde_json::Map<String, Value> {
|
||||
params: Vec<(String, P)>,
|
||||
) -> Map<String, Value>
|
||||
where
|
||||
P: Into<ParamInput>,
|
||||
{
|
||||
let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty());
|
||||
let mut converted = serde_json::Map::with_capacity(params.len());
|
||||
let mut converted = Map::with_capacity(params.len());
|
||||
for (name, value) in params {
|
||||
let value = tool_schema.convert(&name, &value);
|
||||
let value = tool_schema.convert(&name, value.into());
|
||||
converted.insert(name, value);
|
||||
}
|
||||
converted
|
||||
}
|
||||
|
||||
/// Convert one raw string parameter value for one named tool.
|
||||
pub(super) fn convert_param_with_schema(
|
||||
/// Convert one parameter value for one named tool.
|
||||
pub(super) fn convert_param_with_schema<P>(
|
||||
&self,
|
||||
function_name: &str,
|
||||
name: &str,
|
||||
value: &str,
|
||||
) -> Value {
|
||||
value: P,
|
||||
) -> Value
|
||||
where
|
||||
P: Into<ParamInput>,
|
||||
{
|
||||
let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty());
|
||||
tool_schema.convert(name, value)
|
||||
tool_schema.convert(name, value.into())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -101,21 +135,13 @@ impl ToolSchema {
|
||||
Self { params }
|
||||
}
|
||||
|
||||
/// Convert one raw parameter value using its normalized schema type.
|
||||
/// Convert one parameter value using its normalized schema type.
|
||||
///
|
||||
/// If the parameter name is unknown, or we don't have a schema for it, or
|
||||
/// the value fails to convert, this falls back to returning the raw
|
||||
/// string as a JSON string value.
|
||||
fn convert(&self, name: &str, value: &str) -> Value {
|
||||
if value.eq_ignore_ascii_case("null") {
|
||||
return Value::Null;
|
||||
}
|
||||
|
||||
let Some(param_type) = self.params.get(name) else {
|
||||
return Value::String(value.to_string());
|
||||
};
|
||||
|
||||
convert_value(param_type, value).unwrap_or_else(|| Value::String(value.to_string()))
|
||||
/// string as a JSON string value, or object-like JSON for structured input.
|
||||
fn convert(&self, name: &str, input: ParamInput) -> Value {
|
||||
convert_with_optional_schema(self.params.get(name), &input)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -125,7 +151,7 @@ impl JsonParamType {
|
||||
let schema = schema.as_object()?;
|
||||
|
||||
if let Some(type_value) = schema.get("type") {
|
||||
return Self::from_type_value(type_value);
|
||||
return Self::from_type_value(type_value, schema);
|
||||
}
|
||||
|
||||
if let Some(composite) = schema.get("anyOf").or_else(|| schema.get("oneOf")) {
|
||||
@@ -134,32 +160,34 @@ impl JsonParamType {
|
||||
.map(|schemas| schemas.iter().filter_map(Self::from_schema).collect::<Vec<_>>())
|
||||
.filter(|types| !types.is_empty())
|
||||
.map(Self::one_of)
|
||||
.unwrap_or(Self::Object);
|
||||
.unwrap_or_else(|| Self::object_from_schema(Some(schema)));
|
||||
return Some(param_type);
|
||||
}
|
||||
|
||||
// Typically, these types are already handled by checking the "type" field, but
|
||||
// we can also infer them from their characteristic fields if "type" is missing.
|
||||
if schema.contains_key("enum") {
|
||||
return Some(Self::String);
|
||||
}
|
||||
if schema.contains_key("items") {
|
||||
return Some(Self::Array);
|
||||
return Some(Self::array_from_schema(Some(schema)));
|
||||
}
|
||||
if schema.contains_key("properties") {
|
||||
return Some(Self::Object);
|
||||
if schema.contains_key("properties") || schema.contains_key("additionalProperties") {
|
||||
return Some(Self::object_from_schema(Some(schema)));
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// Normalize a JSON schema `type` value.
|
||||
fn from_type_value(type_value: &Value) -> Option<Self> {
|
||||
fn from_type_value(type_value: &Value, schema: &Map<String, Value>) -> Option<Self> {
|
||||
match type_value {
|
||||
Value::String(kind) => Self::from_type_name(kind),
|
||||
Value::String(kind) => Self::from_type_name(kind, Some(schema)),
|
||||
Value::Array(kinds) => {
|
||||
let types = kinds
|
||||
.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.filter_map(Self::from_type_name)
|
||||
.filter_map(|kind| Self::from_type_name(kind, Some(schema)))
|
||||
.collect::<Vec<_>>();
|
||||
if types.is_empty() {
|
||||
None
|
||||
@@ -172,15 +200,15 @@ impl JsonParamType {
|
||||
}
|
||||
|
||||
/// Normalize one JSON schema type name.
|
||||
fn from_type_name(kind: &str) -> Option<Self> {
|
||||
fn from_type_name(kind: &str, schema: Option<&Map<String, Value>>) -> Option<Self> {
|
||||
let kind = kind.trim().to_ascii_lowercase();
|
||||
match kind.as_str() {
|
||||
"string" | "str" | "text" | "varchar" | "char" | "enum" => Some(Self::String),
|
||||
"integer" | "int" => Some(Self::Integer),
|
||||
"number" | "float" | "double" => Some(Self::Number),
|
||||
"boolean" | "bool" | "binary" => Some(Self::Boolean),
|
||||
"object" | "dict" | "map" => Some(Self::Object),
|
||||
"array" | "arr" | "list" | "sequence" => Some(Self::Array),
|
||||
"object" | "dict" | "map" => Some(Self::object_from_schema(schema)),
|
||||
"array" | "arr" | "list" | "sequence" => Some(Self::array_from_schema(schema)),
|
||||
"null" => Some(Self::Null),
|
||||
_ if kind.starts_with("int")
|
||||
|| kind.starts_with("uint")
|
||||
@@ -191,12 +219,52 @@ impl JsonParamType {
|
||||
Some(Self::Integer)
|
||||
}
|
||||
_ if kind.starts_with("num") || kind.starts_with("float") => Some(Self::Number),
|
||||
_ if kind.starts_with("dict") => Some(Self::Object),
|
||||
_ if kind.starts_with("list") => Some(Self::Array),
|
||||
_ if kind.starts_with("dict") => Some(Self::object_from_schema(schema)),
|
||||
_ if kind.starts_with("list") => Some(Self::array_from_schema(schema)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Normalize object schema fields.
|
||||
fn object_from_schema(schema: Option<&Map<String, Value>>) -> Self {
|
||||
let properties = schema
|
||||
.and_then(|schema| schema.get("properties"))
|
||||
.and_then(Value::as_object)
|
||||
.map(|properties| {
|
||||
properties
|
||||
.iter()
|
||||
.filter_map(|(name, schema)| {
|
||||
Self::from_schema(schema).map(|param_type| (name.clone(), param_type))
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let additional_properties =
|
||||
schema.and_then(|schema| schema.get("additionalProperties")).and_then(|schema| {
|
||||
if schema.is_object() {
|
||||
Self::from_schema(schema).map(Box::new)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
});
|
||||
|
||||
Self::Object {
|
||||
properties,
|
||||
additional_properties,
|
||||
}
|
||||
}
|
||||
|
||||
/// Normalize array schema fields.
|
||||
fn array_from_schema(schema: Option<&Map<String, Value>>) -> Self {
|
||||
let items = schema
|
||||
.and_then(|schema| schema.get("items"))
|
||||
.and_then(Self::from_schema)
|
||||
.map(Box::new);
|
||||
|
||||
Self::Array { items }
|
||||
}
|
||||
|
||||
/// Collapse a candidate type list into one normalized type.
|
||||
fn one_of(mut types: Vec<Self>) -> Self {
|
||||
if types.len() == 1 {
|
||||
@@ -207,23 +275,126 @@ impl JsonParamType {
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert one raw string value to a normalized JSON type.
|
||||
fn convert_value(param_type: &JsonParamType, value: &str) -> Option<Value> {
|
||||
match param_type {
|
||||
JsonParamType::String => Some(Value::String(value.to_string())),
|
||||
JsonParamType::Integer => value.parse::<i64>().ok().map(Number::from).map(Value::Number),
|
||||
JsonParamType::Number => convert_number(value),
|
||||
JsonParamType::Boolean => convert_boolean(value),
|
||||
JsonParamType::Object | JsonParamType::Array => serde_json::from_str(value).ok(),
|
||||
JsonParamType::Null => value.eq_ignore_ascii_case("null").then_some(Value::Null),
|
||||
JsonParamType::OneOf(types) => {
|
||||
types.iter().find_map(|param_type| convert_value(param_type, value))
|
||||
/// Convert one parameter input to a normalized JSON value.
|
||||
fn convert_with_optional_schema(param_type: Option<&JsonParamType>, input: &ParamInput) -> Value {
|
||||
// For literal `null`, always convert to JSON null value.
|
||||
if let ParamInput::Text(value) = input
|
||||
&& value.eq_ignore_ascii_case("null")
|
||||
{
|
||||
return Value::Null;
|
||||
}
|
||||
|
||||
// If we have a schema, try to convert the value using it.
|
||||
if let Some(param_type) = param_type
|
||||
&& let Some(value) = try_convert_value(param_type, input)
|
||||
{
|
||||
return value;
|
||||
}
|
||||
// We don't have a schema, or conversion failed, use fallback logic.
|
||||
match input {
|
||||
ParamInput::Text(value) => Value::String(value.clone()),
|
||||
ParamInput::Elements(elements) => {
|
||||
// Convert structured input to object without a schema.
|
||||
Value::Object(convert_elements_to_object(elements, &BTreeMap::new(), None))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert one parameter input to a normalized JSON type.
|
||||
fn try_convert_value(param_type: &JsonParamType, input: &ParamInput) -> Option<Value> {
|
||||
match input {
|
||||
ParamInput::Text(value) => try_convert_text_value(param_type, value),
|
||||
ParamInput::Elements(elements) => try_convert_elements_value(param_type, elements),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert one raw string value to a normalized JSON type.
|
||||
fn try_convert_text_value(param_type: &JsonParamType, value: &str) -> Option<Value> {
|
||||
match param_type {
|
||||
JsonParamType::String => Some(Value::String(value.to_string())),
|
||||
JsonParamType::Integer => value.parse::<i64>().ok().map(Number::from).map(Value::Number),
|
||||
JsonParamType::Number => try_convert_number(value),
|
||||
JsonParamType::Boolean => try_convert_boolean(value),
|
||||
JsonParamType::Object { .. } if value.is_empty() => Some(Value::Object(Map::new())),
|
||||
JsonParamType::Array { .. } if value.is_empty() => Some(Value::Array(Vec::new())),
|
||||
JsonParamType::Object { .. } | JsonParamType::Array { .. } => {
|
||||
// For composite types with string input, simply interpret the string as JSON.
|
||||
serde_json::from_str(value).ok()
|
||||
}
|
||||
JsonParamType::Null => value.eq_ignore_ascii_case("null").then_some(Value::Null),
|
||||
JsonParamType::OneOf(types) => {
|
||||
types.iter().find_map(|param_type| try_convert_text_value(param_type, value))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert one structured parameter input to a normalized JSON type.
|
||||
fn try_convert_elements_value(
|
||||
param_type: &JsonParamType,
|
||||
elements: &[ParamElement],
|
||||
) -> Option<Value> {
|
||||
match param_type {
|
||||
JsonParamType::Object {
|
||||
properties,
|
||||
additional_properties,
|
||||
} => Some(Value::Object(convert_elements_to_object(
|
||||
elements,
|
||||
properties,
|
||||
additional_properties.as_deref(),
|
||||
))),
|
||||
JsonParamType::Array { items } => Some(Value::Array(
|
||||
// Collect all child elements into an array, regardless of their names.
|
||||
elements
|
||||
.iter()
|
||||
.map(|element| convert_with_optional_schema(items.as_deref(), &element.value))
|
||||
.collect(),
|
||||
)),
|
||||
JsonParamType::OneOf(types) => types
|
||||
.iter()
|
||||
.find_map(|param_type| try_convert_elements_value(param_type, elements)),
|
||||
|
||||
// Primitive types can't be converted from structured input.
|
||||
JsonParamType::String
|
||||
| JsonParamType::Integer
|
||||
| JsonParamType::Number
|
||||
| JsonParamType::Boolean
|
||||
| JsonParamType::Null => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert structured elements to an object, using field schemas when present.
|
||||
fn convert_elements_to_object(
|
||||
elements: &[ParamElement],
|
||||
properties: &BTreeMap<String, JsonParamType>,
|
||||
additional_properties: Option<&JsonParamType>,
|
||||
) -> Map<String, Value> {
|
||||
let mut object = Map::with_capacity(elements.len());
|
||||
for element in elements {
|
||||
let param_type = properties.get(&element.name).or(additional_properties);
|
||||
let value = convert_with_optional_schema(param_type, &element.value);
|
||||
insert_object_value(&mut object, element.name.clone(), value);
|
||||
}
|
||||
object
|
||||
}
|
||||
|
||||
/// Insert an object field while preserving duplicate keys as arrays.
|
||||
fn insert_object_value(object: &mut Map<String, Value>, key: String, value: Value) {
|
||||
if let Some(existing) = object.get_mut(&key) {
|
||||
match existing {
|
||||
// Collect values under the same key into an array.
|
||||
Value::Array(values) => values.push(value),
|
||||
existing => {
|
||||
let first = std::mem::replace(existing, Value::Null);
|
||||
*existing = Value::Array(vec![first, value]);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
object.insert(key, value);
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert one raw string value to a JSON number.
|
||||
fn convert_number(value: &str) -> Option<Value> {
|
||||
fn try_convert_number(value: &str) -> Option<Value> {
|
||||
serde_json::from_str::<Number>(value)
|
||||
.or_else(|_| value.parse::<i64>().map(Number::from))
|
||||
.or_else(|_| value.parse::<f64>().ok().and_then(Number::from_f64).ok_or(()))
|
||||
@@ -232,7 +403,7 @@ fn convert_number(value: &str) -> Option<Value> {
|
||||
}
|
||||
|
||||
/// Convert one raw string value to a boolean.
|
||||
fn convert_boolean(value: &str) -> Option<Value> {
|
||||
fn try_convert_boolean(value: &str) -> Option<Value> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"true" | "1" => Some(Value::Bool(true)),
|
||||
"false" | "0" => Some(Value::Bool(false)),
|
||||
@@ -242,9 +413,9 @@ fn convert_boolean(value: &str) -> Option<Value> {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::{ToolSchema, ToolSchemas};
|
||||
use super::{ParamElement, ParamInput, ToolSchema, ToolSchemas};
|
||||
use crate::Tool;
|
||||
|
||||
fn test_tool(name: &str, parameters: serde_json::Value) -> Tool {
|
||||
@@ -260,8 +431,8 @@ mod tests {
|
||||
fn invalid_schema_converts_everything_as_string() {
|
||||
let params = ToolSchema::from_schema(&json!({ "type": "object" }));
|
||||
|
||||
assert_eq!(params.convert("count", "42"), json!("42"));
|
||||
assert_eq!(params.convert("count", "null"), json!(null));
|
||||
assert_eq!(params.convert("count", text("42")), json!("42"));
|
||||
assert_eq!(params.convert("count", text("null")), json!(null));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -275,9 +446,9 @@ mod tests {
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(params.convert("unknown_schema", "42"), json!("42"));
|
||||
assert_eq!(params.convert("unknown_type", "42"), json!("42"));
|
||||
assert_eq!(params.convert("known", "42"), json!(42));
|
||||
assert_eq!(params.convert("unknown_schema", text("42")), json!("42"));
|
||||
assert_eq!(params.convert("unknown_type", text("42")), json!("42"));
|
||||
assert_eq!(params.convert("known", text("42")), json!(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -298,16 +469,25 @@ mod tests {
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(params.convert("text", "42"), json!("42"));
|
||||
assert_eq!(params.convert("count", "42"), json!(42));
|
||||
assert_eq!(params.convert("size", "5.0"), json!(5.0));
|
||||
assert_eq!(params.convert("ratio", "2.5"), json!(2.5));
|
||||
assert_eq!(params.convert("enabled", "1"), json!(true));
|
||||
assert_eq!(params.convert("payload", r#"{"k":1}"#), json!({ "k": 1 }));
|
||||
assert_eq!(params.convert("mapping", r#"{"k":1}"#), json!({ "k": 1 }));
|
||||
assert_eq!(params.convert("items", "[1,2]"), json!([1, 2]));
|
||||
assert_eq!(params.convert("names", r#"["a","b"]"#), json!(["a", "b"]));
|
||||
assert_eq!(params.convert("nothing", "null"), json!(null));
|
||||
assert_eq!(params.convert("text", text("42")), json!("42"));
|
||||
assert_eq!(params.convert("count", text("42")), json!(42));
|
||||
assert_eq!(params.convert("size", text("5.0")), json!(5.0));
|
||||
assert_eq!(params.convert("ratio", text("2.5")), json!(2.5));
|
||||
assert_eq!(params.convert("enabled", text("1")), json!(true));
|
||||
assert_eq!(
|
||||
params.convert("payload", text(r#"{"k":1}"#)),
|
||||
json!({ "k": 1 })
|
||||
);
|
||||
assert_eq!(
|
||||
params.convert("mapping", text(r#"{"k":1}"#)),
|
||||
json!({ "k": 1 })
|
||||
);
|
||||
assert_eq!(params.convert("items", text("[1,2]")), json!([1, 2]));
|
||||
assert_eq!(
|
||||
params.convert("names", text(r#"["a","b"]"#)),
|
||||
json!(["a", "b"])
|
||||
);
|
||||
assert_eq!(params.convert("nothing", text("null")), json!(null));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -333,7 +513,22 @@ mod tests {
|
||||
}
|
||||
|
||||
fn converted_number_text(params: &ToolSchema, value: &str) -> String {
|
||||
serde_json::to_string(¶ms.convert("value", value)).unwrap()
|
||||
serde_json::to_string(¶ms.convert("value", text(value))).unwrap()
|
||||
}
|
||||
|
||||
fn text(value: &str) -> ParamInput {
|
||||
ParamInput::Text(value.to_string())
|
||||
}
|
||||
|
||||
fn elem(name: &str, value: ParamInput) -> ParamElement {
|
||||
ParamElement {
|
||||
name: name.to_string(),
|
||||
value,
|
||||
}
|
||||
}
|
||||
|
||||
fn elements(elements: Vec<ParamElement>) -> ParamInput {
|
||||
ParamInput::Elements(elements)
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -350,12 +545,12 @@ mod tests {
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(params.convert("s", "x"), json!("x"));
|
||||
assert_eq!(params.convert("i", "7"), json!(7));
|
||||
assert_eq!(params.convert("n", "7.5"), json!(7.5));
|
||||
assert_eq!(params.convert("b", "true"), json!(true));
|
||||
assert_eq!(params.convert("a", "[1]"), json!([1]));
|
||||
assert_eq!(params.convert("o", r#"{"x":1}"#), json!({ "x": 1 }));
|
||||
assert_eq!(params.convert("s", text("x")), json!("x"));
|
||||
assert_eq!(params.convert("i", text("7")), json!(7));
|
||||
assert_eq!(params.convert("n", text("7.5")), json!(7.5));
|
||||
assert_eq!(params.convert("b", text("true")), json!(true));
|
||||
assert_eq!(params.convert("a", text("[1]")), json!([1]));
|
||||
assert_eq!(params.convert("o", text(r#"{"x":1}"#)), json!({ "x": 1 }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -373,8 +568,8 @@ mod tests {
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(integer_first.convert("value", "42"), json!(42));
|
||||
assert_eq!(string_first.convert("value", "42"), json!("42"));
|
||||
assert_eq!(integer_first.convert("value", text("42")), json!(42));
|
||||
assert_eq!(string_first.convert("value", text("42")), json!("42"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -396,9 +591,9 @@ mod tests {
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(params.convert("choice", "42"), json!(42));
|
||||
assert_eq!(params.convert("choice", text("42")), json!(42));
|
||||
assert_eq!(
|
||||
params.convert("fallback_object", r#"{"x":1}"#),
|
||||
params.convert("fallback_object", text(r#"{"x":1}"#)),
|
||||
json!({ "x": 1 })
|
||||
);
|
||||
}
|
||||
@@ -414,9 +609,12 @@ mod tests {
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(params.convert("choice", "a"), json!("a"));
|
||||
assert_eq!(params.convert("items", "[1,2]"), json!([1, 2]));
|
||||
assert_eq!(params.convert("payload", r#"{"x":1}"#), json!({ "x": 1 }));
|
||||
assert_eq!(params.convert("choice", text("a")), json!("a"));
|
||||
assert_eq!(params.convert("items", text("[1,2]")), json!([1, 2]));
|
||||
assert_eq!(
|
||||
params.convert("payload", text(r#"{"x":1}"#)),
|
||||
json!({ "x": 1 })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -518,4 +716,162 @@ mod tests {
|
||||
assert_eq!(converted.get("topn"), Some(&json!("5")));
|
||||
assert_eq!(converted.get("nullish"), Some(&json!(null)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_structured_inputs_with_recursive_schema() {
|
||||
let schemas = ToolSchemas::from_tools(&[test_tool(
|
||||
"create_order",
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"user_id": { "type": "integer" },
|
||||
"urgent": { "type": "boolean" },
|
||||
"note": { "type": "string" },
|
||||
"nil": { "type": "string" },
|
||||
"shipping": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": { "type": "string" },
|
||||
"zip": { "type": "integer" }
|
||||
}
|
||||
},
|
||||
"items": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"sku": { "type": "string" },
|
||||
"qty": { "type": "integer" }
|
||||
}
|
||||
}
|
||||
},
|
||||
"metadata": {
|
||||
"type": "object",
|
||||
"additionalProperties": { "type": "integer" }
|
||||
},
|
||||
"duplicate_demo": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"tag": { "type": "string" }
|
||||
}
|
||||
},
|
||||
"schema_mismatch_array": {
|
||||
"type": "array",
|
||||
"items": { "type": "integer" }
|
||||
},
|
||||
"closed_object": {
|
||||
"type": "object",
|
||||
"additionalProperties": false
|
||||
},
|
||||
"open_object": {
|
||||
"type": "object",
|
||||
"additionalProperties": true
|
||||
},
|
||||
"payload_text": { "type": "object" },
|
||||
"items_text": { "type": "array" }
|
||||
}
|
||||
}),
|
||||
)]);
|
||||
|
||||
let converted = schemas.convert_params_with_schema(
|
||||
"create_order",
|
||||
vec![
|
||||
("user_id".to_string(), text("42")),
|
||||
("urgent".to_string(), text("true")),
|
||||
("note".to_string(), text("Please leave at front desk.")),
|
||||
("nil".to_string(), text("NULL")),
|
||||
(
|
||||
"shipping".to_string(),
|
||||
elements(vec![
|
||||
elem("city", text("Singapore")),
|
||||
elem("zip", text("018956")),
|
||||
]),
|
||||
),
|
||||
(
|
||||
"items".to_string(),
|
||||
elements(vec![
|
||||
elem(
|
||||
"item1",
|
||||
elements(vec![elem("sku", text("book-001")), elem("qty", text("2"))]),
|
||||
),
|
||||
elem(
|
||||
"item2",
|
||||
elements(vec![elem("sku", text("pen-007")), elem("qty", text("5"))]),
|
||||
),
|
||||
]),
|
||||
),
|
||||
(
|
||||
"metadata".to_string(),
|
||||
elements(vec![elem("score", text("42")), elem("rank", text("7"))]),
|
||||
),
|
||||
(
|
||||
"duplicate_demo".to_string(),
|
||||
elements(vec![elem("tag", text("a")), elem("tag", text("b"))]),
|
||||
),
|
||||
(
|
||||
"closed_object".to_string(),
|
||||
elements(vec![elem("unknown", text("x"))]),
|
||||
),
|
||||
(
|
||||
"open_object".to_string(),
|
||||
elements(vec![elem("unknown", text("y"))]),
|
||||
),
|
||||
("payload_text".to_string(), text(r#"{"x":1}"#)),
|
||||
("items_text".to_string(), text("[1,2]")),
|
||||
(
|
||||
"unknown_struct".to_string(),
|
||||
elements(vec![
|
||||
elem("a", text("1")),
|
||||
elem("a", text("2")),
|
||||
elem("nil", text("null")),
|
||||
]),
|
||||
),
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
Value::Object(converted),
|
||||
json!({
|
||||
"user_id": 42,
|
||||
"urgent": true,
|
||||
"note": "Please leave at front desk.",
|
||||
"nil": null,
|
||||
"shipping": {
|
||||
"city": "Singapore",
|
||||
"zip": 18956
|
||||
},
|
||||
"items": [
|
||||
{
|
||||
"sku": "book-001",
|
||||
"qty": 2
|
||||
},
|
||||
{
|
||||
"sku": "pen-007",
|
||||
"qty": 5
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"score": 42,
|
||||
"rank": 7
|
||||
},
|
||||
"duplicate_demo": {
|
||||
"tag": ["a", "b"]
|
||||
},
|
||||
"closed_object": {
|
||||
"unknown": "x"
|
||||
},
|
||||
"open_object": {
|
||||
"unknown": "y"
|
||||
},
|
||||
"payload_text": {
|
||||
"x": 1
|
||||
},
|
||||
"items_text": [1, 2],
|
||||
"unknown_struct": {
|
||||
"a": ["1", "2"],
|
||||
"nil": null
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,11 +2,15 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import json
|
||||
from argparse import Namespace
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pybase64 as base64
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
import vllm.benchmarks.datasets.datasets as datasets_module
|
||||
from vllm.benchmarks.datasets import CustomImageDataset, get_samples
|
||||
from vllm.benchmarks.lib.endpoint_request_func import (
|
||||
RequestFuncInput,
|
||||
@@ -33,6 +37,22 @@ def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
|
||||
f.write(json.dumps(row) + "\n")
|
||||
|
||||
|
||||
def _write_png(path: Path, color: tuple[int, int, int] = (255, 0, 0)) -> None:
|
||||
Image.new("RGB", (1, 1), color=color).save(path)
|
||||
|
||||
|
||||
def _decode_data_url(data_url: str) -> tuple[str, bytes]:
|
||||
prefix, image_base64 = data_url.split(",", 1)
|
||||
return prefix, base64.b64decode(image_base64)
|
||||
|
||||
|
||||
def _assert_png_data_url(data_url: str) -> None:
|
||||
prefix, image_bytes = _decode_data_url(data_url)
|
||||
assert prefix == "data:image/png;base64"
|
||||
with Image.open(BytesIO(image_bytes)) as image:
|
||||
image.verify()
|
||||
|
||||
|
||||
def _args_for_custom_image(dataset_path: Path) -> Namespace:
|
||||
return Namespace(
|
||||
dataset_name="custom_image",
|
||||
@@ -42,6 +62,7 @@ def _args_for_custom_image(dataset_path: Path) -> Namespace:
|
||||
num_prompts=2,
|
||||
custom_output_len=32,
|
||||
enable_multimodal_chat=False,
|
||||
custom_ensure_client_side_data=False,
|
||||
request_id_prefix="req-",
|
||||
no_oversample=False,
|
||||
)
|
||||
@@ -230,6 +251,125 @@ def test_custom_image_dataset_wraps_interleaved_content_for_multimodal_chat(
|
||||
assert _get_chat_messages(request_input) == sample.prompt
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_custom_image_dataset_encodes_image_media_when_requested(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
image_a = tmp_path / "chart_a.png"
|
||||
image_b = tmp_path / "chart b.png"
|
||||
_write_png(image_a, color=(255, 0, 0))
|
||||
_write_png(image_b, color=(0, 255, 0))
|
||||
data_url = "data:image/png;base64,Zm9v"
|
||||
remote_url = "https://example.com/chart.png"
|
||||
original_fetch_image = datasets_module.fetch_image
|
||||
|
||||
def fake_fetch_image(image_url: str) -> Image.Image:
|
||||
if image_url == remote_url:
|
||||
return Image.new("RGB", (1, 1), color=(0, 0, 255))
|
||||
return original_fetch_image(image_url)
|
||||
|
||||
monkeypatch.setattr(datasets_module, "fetch_image", fake_fetch_image)
|
||||
|
||||
jsonl = tmp_path / "images.jsonl"
|
||||
_write_jsonl(
|
||||
jsonl,
|
||||
[
|
||||
{
|
||||
"prompt": "Compare the charts.",
|
||||
"image_files": [
|
||||
str(image_a),
|
||||
image_b.as_uri(),
|
||||
remote_url,
|
||||
data_url,
|
||||
],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
|
||||
samples = dataset.sample(
|
||||
tokenizer=_Tokenizer(),
|
||||
num_requests=1,
|
||||
output_len=32,
|
||||
ensure_client_side_data=True,
|
||||
)
|
||||
|
||||
assert len(samples) == 1
|
||||
assert isinstance(samples[0].multi_modal_data, list)
|
||||
image_urls = [part["image_url"]["url"] for part in samples[0].multi_modal_data]
|
||||
|
||||
_assert_png_data_url(image_urls[0])
|
||||
_assert_png_data_url(image_urls[1])
|
||||
_assert_png_data_url(image_urls[2])
|
||||
assert image_urls[3] == data_url
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_custom_image_dataset_encodes_interleaved_image_media(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
image_a = tmp_path / "chart_a.png"
|
||||
image_b = tmp_path / "chart_b.png"
|
||||
_write_png(image_a, color=(255, 0, 0))
|
||||
_write_png(image_b, color=(0, 255, 0))
|
||||
jsonl = tmp_path / "images.jsonl"
|
||||
_write_jsonl(
|
||||
jsonl,
|
||||
[
|
||||
{
|
||||
"content": [
|
||||
{"type": "text", "text": "Compare "},
|
||||
{"type": "image", "image": str(image_a)},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": image_b.as_uri(),
|
||||
"detail": "low",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
|
||||
samples = dataset.sample(
|
||||
tokenizer=_Tokenizer(),
|
||||
num_requests=1,
|
||||
output_len=32,
|
||||
ensure_client_side_data=True,
|
||||
)
|
||||
|
||||
sample = samples[0]
|
||||
assert isinstance(sample.prompt, list)
|
||||
_assert_png_data_url(sample.prompt[1]["image_url"]["url"])
|
||||
_assert_png_data_url(sample.prompt[2]["image_url"]["url"])
|
||||
assert sample.prompt[2]["image_url"]["detail"] == "low"
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_custom_image_dataset_rejects_invalid_image_media(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
invalid_image = tmp_path / "not_an_image.png"
|
||||
invalid_image.write_text("not an image")
|
||||
jsonl = tmp_path / "images.jsonl"
|
||||
_write_jsonl(
|
||||
jsonl,
|
||||
[{"prompt": "Describe the image.", "image_files": [str(invalid_image)]}],
|
||||
)
|
||||
|
||||
dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
|
||||
with pytest.raises(ValueError, match="Invalid image URL"):
|
||||
dataset.sample(
|
||||
tokenizer=_Tokenizer(),
|
||||
num_requests=1,
|
||||
output_len=32,
|
||||
ensure_client_side_data=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_custom_image_dataset_rejects_invalid_content_part(
|
||||
tmp_path: Path,
|
||||
|
||||
@@ -55,6 +55,15 @@ def test_dynamic_shapes_compilation(
|
||||
evaluate_guards,
|
||||
):
|
||||
"""Test that all dynamic shapes types compile successfully"""
|
||||
if shapes_type == DynamicShapesType.UNBACKED and not is_torch_equal_or_newer(
|
||||
"2.11.0"
|
||||
):
|
||||
# NOTE[ROCm]: shape_id (used by Qwen2/Llama to relate input dims) only
|
||||
# landed in torch 2.11, but the ROCm CI still runs torch 2.10.x. On
|
||||
# older torch there's no way to express it, so unbacked shapes go
|
||||
# data-dependent and compilation blows up -- nothing to test.
|
||||
pytest.skip("unbacked dynamic shapes with shape_id require torch>=2.11")
|
||||
|
||||
if evaluate_guards and shapes_type == DynamicShapesType.UNBACKED:
|
||||
pytest.skip("unbacked dynamic shapes do not add guards")
|
||||
|
||||
|
||||
@@ -284,7 +284,7 @@ def test_multiproc_executor_pipeline_parallel():
|
||||
assert output_rank == 2, "Output rank should be 2 (first rank of last PP stage)"
|
||||
|
||||
# Verify max_concurrent_batches for pipeline parallel
|
||||
assert executor.max_concurrent_batches == 2, (
|
||||
assert vllm_config.max_concurrent_batches == 2, (
|
||||
"Max concurrent batches should equal PP size"
|
||||
)
|
||||
|
||||
|
||||
@@ -83,7 +83,7 @@ def assert_executor(executor, tp_size, pp_size):
|
||||
assert executor._get_output_rank() == expected_output_rank
|
||||
|
||||
if pp_size > 1:
|
||||
assert executor.max_concurrent_batches == pp_size
|
||||
assert executor.vllm_config.max_concurrent_batches == pp_size
|
||||
|
||||
executor.check_health()
|
||||
assert not executor.is_failed
|
||||
|
||||
@@ -18,6 +18,7 @@ from torch.multiprocessing.reductions import reduce_tensor
|
||||
from vllm.config.parallel import ParallelConfig
|
||||
from vllm.config.weight_transfer import WeightTransferConfig
|
||||
from vllm.distributed.weight_transfer import WeightTransferEngineFactory
|
||||
from vllm.distributed.weight_transfer.base import SparseWeightPatch
|
||||
from vllm.distributed.weight_transfer.ipc_engine import (
|
||||
IPCWeightTransferEngine,
|
||||
IPCWeightTransferInitInfo,
|
||||
@@ -89,6 +90,67 @@ class TestNCCLWeightTransferUpdateInfoValidation:
|
||||
)
|
||||
assert len(info.names) == 0
|
||||
|
||||
def test_valid_sparse_update_info(self):
|
||||
"""Test creating valid sparse NCCL update info."""
|
||||
info = NCCLWeightTransferUpdateInfo(
|
||||
names=["layer.weight", "layer.bias"],
|
||||
dtype_names=["float32", "bfloat16"],
|
||||
shapes=[[10, 10], [10]],
|
||||
num_updates_list=[4, 2],
|
||||
update_kind="sparse_flat",
|
||||
)
|
||||
assert info.update_kind == "sparse_flat"
|
||||
assert info.num_updates_list == [4, 2]
|
||||
|
||||
def test_sparse_update_requires_num_updates_list(self):
|
||||
with pytest.raises(ValueError, match="`num_updates_list` is required"):
|
||||
NCCLWeightTransferUpdateInfo(
|
||||
names=["layer.weight"],
|
||||
dtype_names=["float32"],
|
||||
shapes=[[10, 10]],
|
||||
update_kind="sparse_flat",
|
||||
)
|
||||
|
||||
def test_sparse_update_rejects_empty_num_updates_list(self):
|
||||
with pytest.raises(ValueError, match="cannot be empty"):
|
||||
NCCLWeightTransferUpdateInfo(
|
||||
names=[],
|
||||
dtype_names=[],
|
||||
shapes=[],
|
||||
num_updates_list=[],
|
||||
update_kind="sparse_flat",
|
||||
)
|
||||
|
||||
def test_sparse_update_rejects_packed(self):
|
||||
with pytest.raises(ValueError, match="cannot be combined with `packed=True`"):
|
||||
NCCLWeightTransferUpdateInfo(
|
||||
names=["layer.weight"],
|
||||
dtype_names=["float32"],
|
||||
shapes=[[10, 10]],
|
||||
num_updates_list=[3],
|
||||
update_kind="sparse_flat",
|
||||
packed=True,
|
||||
)
|
||||
|
||||
def test_sparse_update_rejects_mismatched_num_updates(self):
|
||||
with pytest.raises(ValueError, match="`num_updates_list`"):
|
||||
NCCLWeightTransferUpdateInfo(
|
||||
names=["layer.weight", "layer.bias"],
|
||||
dtype_names=["float32", "float32"],
|
||||
shapes=[[10, 10], [10]],
|
||||
num_updates_list=[3],
|
||||
update_kind="sparse_flat",
|
||||
)
|
||||
|
||||
def test_dense_update_rejects_sparse_metadata(self):
|
||||
with pytest.raises(ValueError, match="Sparse metadata"):
|
||||
NCCLWeightTransferUpdateInfo(
|
||||
names=["layer.weight"],
|
||||
dtype_names=["float32"],
|
||||
shapes=[[10, 10]],
|
||||
num_updates_list=[3],
|
||||
)
|
||||
|
||||
|
||||
# --- Unit Tests: Engine Parsing ---
|
||||
|
||||
@@ -222,6 +284,27 @@ def test_nccl_receive_weights_without_init_raises():
|
||||
engine.receive_weights(update_info, lambda x: None)
|
||||
|
||||
|
||||
def test_nccl_receive_sparse_weights_without_init_raises():
|
||||
"""Test that sparse receive raises if init_transfer_engine wasn't called."""
|
||||
if torch.accelerator.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU for this test")
|
||||
|
||||
config = WeightTransferConfig(backend="nccl")
|
||||
parallel_config = create_mock_parallel_config()
|
||||
engine = NCCLWeightTransferEngine(config, parallel_config)
|
||||
|
||||
update_info = NCCLWeightTransferUpdateInfo(
|
||||
names=["w"],
|
||||
dtype_names=["float32"],
|
||||
shapes=[[10]],
|
||||
num_updates_list=[2],
|
||||
update_kind="sparse_flat",
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="not initialized"):
|
||||
engine.receive_sparse_weights(update_info, lambda x: None)
|
||||
|
||||
|
||||
# --- Integration Test: NCCL Weight Transfer Between Ray Tasks ---
|
||||
|
||||
|
||||
@@ -379,6 +462,136 @@ def test_nccl_weight_transfer_between_processes():
|
||||
)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
def trainer_broadcast_sparse_tensor(
|
||||
master_address: str,
|
||||
master_port: int,
|
||||
world_size: int,
|
||||
) -> bool:
|
||||
"""Trainer task that broadcasts sparse patches via NCCL."""
|
||||
import torch
|
||||
|
||||
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
|
||||
from vllm.distributed.utils import StatelessProcessGroup
|
||||
from vllm.distributed.weight_transfer.base import SparseWeightPatch
|
||||
from vllm.distributed.weight_transfer.nccl_engine import (
|
||||
NCCLTrainerSendWeightsArgs,
|
||||
NCCLWeightTransferEngine,
|
||||
)
|
||||
|
||||
pg = StatelessProcessGroup.create(
|
||||
host=master_address,
|
||||
port=master_port,
|
||||
rank=0,
|
||||
world_size=world_size,
|
||||
)
|
||||
comm = PyNcclCommunicator(pg, device=0)
|
||||
|
||||
patch = SparseWeightPatch(
|
||||
name="test.weight",
|
||||
indices=torch.tensor([1, 7, 25], dtype=torch.int32, device="cuda:0"),
|
||||
values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device="cuda:0"),
|
||||
)
|
||||
NCCLWeightTransferEngine.trainer_send_sparse_weights(
|
||||
iter([patch]),
|
||||
NCCLTrainerSendWeightsArgs(group=comm),
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
return True
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
def inference_receive_sparse_tensor(
|
||||
master_address: str,
|
||||
master_port: int,
|
||||
world_size: int,
|
||||
) -> dict:
|
||||
"""Inference task that receives sparse patches via NCCLWeightTransferEngine."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config.parallel import ParallelConfig
|
||||
from vllm.config.weight_transfer import WeightTransferConfig
|
||||
from vllm.distributed.weight_transfer.nccl_engine import (
|
||||
NCCLWeightTransferEngine,
|
||||
NCCLWeightTransferInitInfo,
|
||||
NCCLWeightTransferUpdateInfo,
|
||||
)
|
||||
|
||||
config = WeightTransferConfig(backend="nccl")
|
||||
parallel_config = MagicMock(spec=ParallelConfig)
|
||||
parallel_config.rank = 0
|
||||
parallel_config.world_size = 1
|
||||
parallel_config.data_parallel_rank = 0
|
||||
parallel_config.data_parallel_index = 0
|
||||
|
||||
engine = NCCLWeightTransferEngine(config, parallel_config)
|
||||
engine.init_transfer_engine(
|
||||
NCCLWeightTransferInitInfo(
|
||||
master_address=master_address,
|
||||
master_port=master_port,
|
||||
rank_offset=1,
|
||||
world_size=world_size,
|
||||
)
|
||||
)
|
||||
|
||||
target = torch.zeros(30, dtype=torch.float32, device="cuda")
|
||||
|
||||
def apply_sparse_patches(patches: list[SparseWeightPatch]):
|
||||
for patch in patches:
|
||||
target.index_copy_(0, patch.indices.to(torch.long), patch.values)
|
||||
|
||||
update_info = NCCLWeightTransferUpdateInfo(
|
||||
names=["test.weight"],
|
||||
dtype_names=["float32"],
|
||||
shapes=[[30]],
|
||||
num_updates_list=[3],
|
||||
update_kind="sparse_flat",
|
||||
)
|
||||
engine.receive_sparse_weights(update_info, apply_sparse_patches)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
expected = torch.zeros(30, dtype=torch.float32, device="cuda")
|
||||
expected[[1, 7, 25]] = torch.tensor(
|
||||
[10.0, 20.0, 30.0], dtype=torch.float32, device="cuda"
|
||||
)
|
||||
success = torch.equal(target, expected)
|
||||
engine.shutdown()
|
||||
return {
|
||||
"success": success,
|
||||
"selected_values": target[[1, 7, 25]].cpu().tolist(),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 2,
|
||||
reason="Need at least 2 GPUs to run NCCL sparse weight transfer test.",
|
||||
)
|
||||
def test_nccl_sparse_weight_transfer_between_processes():
|
||||
"""Test NCCL sparse weight transfer from trainer to inference process."""
|
||||
ray.init(ignore_reinit_error=True)
|
||||
|
||||
master_address = "127.0.0.1"
|
||||
master_port = get_open_port()
|
||||
world_size = 2
|
||||
|
||||
inference_future = inference_receive_sparse_tensor.remote(
|
||||
master_address, master_port, world_size
|
||||
)
|
||||
trainer_future = trainer_broadcast_sparse_tensor.remote(
|
||||
master_address, master_port, world_size
|
||||
)
|
||||
|
||||
trainer_result, result = ray.get([trainer_future, inference_future])
|
||||
|
||||
assert trainer_result, "Trainer should complete successfully"
|
||||
assert result["success"], (
|
||||
"Sparse weight transfer failed. "
|
||||
f"Received selected values: {result['selected_values']}"
|
||||
)
|
||||
|
||||
|
||||
# --- Unit Tests: IPCWeightTransferUpdateInfo Validation ---
|
||||
|
||||
|
||||
@@ -461,9 +674,101 @@ class TestIPCWeightTransferUpdateInfoValidation:
|
||||
ipc_handles=ipc_handles,
|
||||
)
|
||||
|
||||
def test_missing_ipc_handles_raises(self):
|
||||
"""Test that omitting ipc_handles raises TypeError."""
|
||||
with pytest.raises(TypeError):
|
||||
def test_sparse_update_kind_rejected(self):
|
||||
"""Test that IPC backend rejects sparse update metadata."""
|
||||
if torch.accelerator.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU for this test")
|
||||
|
||||
dummy_tensor = torch.ones(10, 10, device="cuda:0")
|
||||
ipc_handle = reduce_tensor(dummy_tensor)
|
||||
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
||||
ipc_handles = [{gpu_uuid: ipc_handle}]
|
||||
|
||||
with pytest.raises(NotImplementedError, match="dense updates"):
|
||||
IPCWeightTransferUpdateInfo(
|
||||
names=["layer.weight"],
|
||||
dtype_names=["float32"],
|
||||
shapes=[[10, 10]],
|
||||
num_updates_list=[1],
|
||||
ipc_handles=ipc_handles,
|
||||
update_kind="sparse_flat",
|
||||
)
|
||||
|
||||
def test_sparse_methods_not_supported(self):
|
||||
"""Test that IPC engine inherits sparse rejection from the base class."""
|
||||
config = WeightTransferConfig(backend="ipc")
|
||||
parallel_config = create_mock_parallel_config()
|
||||
engine = IPCWeightTransferEngine(
|
||||
config, parallel_config, MagicMock(spec=torch.nn.Module)
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"):
|
||||
engine.receive_sparse_weights(MagicMock(), lambda _: None)
|
||||
with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"):
|
||||
engine.trainer_send_sparse_weights(
|
||||
iter([]),
|
||||
{"mode": "http", "url": "http://localhost:8000"},
|
||||
)
|
||||
|
||||
def test_valid_update_info_from_pickled(self, monkeypatch):
|
||||
"""Test creating IPCWeightTransferUpdateInfo from pickled handles."""
|
||||
if torch.accelerator.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU for this test")
|
||||
|
||||
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
|
||||
dummy_tensor = torch.ones(10, 10, device="cuda:0")
|
||||
ipc_handle = reduce_tensor(dummy_tensor)
|
||||
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
||||
ipc_handles = [{gpu_uuid: ipc_handle}]
|
||||
|
||||
pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
|
||||
|
||||
info = IPCWeightTransferUpdateInfo(
|
||||
names=["layer.weight"],
|
||||
dtype_names=["float32"],
|
||||
shapes=[[10, 10]],
|
||||
ipc_handles_pickled=pickled,
|
||||
)
|
||||
assert info.ipc_handles == ipc_handles
|
||||
assert info.ipc_handles_pickled is None
|
||||
|
||||
def test_pickled_requires_insecure_serialization_flag(self, monkeypatch):
|
||||
"""Test that pickled handles are rejected unless env flag is enabled."""
|
||||
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0")
|
||||
|
||||
with pytest.raises(ValueError, match="VLLM_ALLOW_INSECURE_SERIALIZATION=1"):
|
||||
IPCWeightTransferUpdateInfo(
|
||||
names=[],
|
||||
dtype_names=[],
|
||||
shapes=[],
|
||||
ipc_handles_pickled=base64.b64encode(pickle.dumps([])).decode("utf-8"),
|
||||
)
|
||||
|
||||
def test_both_handles_and_pickled_raises(self):
|
||||
"""Test that providing both ipc_handles and ipc_handles_pickled raises."""
|
||||
if torch.accelerator.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU for this test")
|
||||
|
||||
dummy_tensor = torch.ones(10, 10, device="cuda:0")
|
||||
ipc_handle = reduce_tensor(dummy_tensor)
|
||||
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
||||
ipc_handles = [{gpu_uuid: ipc_handle}]
|
||||
|
||||
pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match="Cannot specify both"):
|
||||
IPCWeightTransferUpdateInfo(
|
||||
names=["layer.weight"],
|
||||
dtype_names=["float32"],
|
||||
shapes=[[10, 10]],
|
||||
ipc_handles=ipc_handles,
|
||||
ipc_handles_pickled=pickled,
|
||||
)
|
||||
|
||||
def test_neither_handles_nor_pickled_raises(self):
|
||||
"""Test that providing neither ipc_handles nor ipc_handles_pickled raises."""
|
||||
with pytest.raises(ValueError, match="must be provided"):
|
||||
IPCWeightTransferUpdateInfo(
|
||||
names=["layer.weight"],
|
||||
dtype_names=["float32"],
|
||||
@@ -558,6 +863,28 @@ class TestIPCEngineParsing:
|
||||
assert gpu_uuid in update_info.ipc_handles[0]
|
||||
assert gpu_uuid in update_info.ipc_handles[1]
|
||||
|
||||
def test_parse_update_info_ignores_none_pickled_handles(self):
|
||||
"""Test Ray/asdict payloads with a null pickled field use ipc_handles."""
|
||||
config = WeightTransferConfig(backend="ipc")
|
||||
parallel_config = create_mock_parallel_config()
|
||||
engine = IPCWeightTransferEngine(
|
||||
config, parallel_config, MagicMock(spec=torch.nn.Module)
|
||||
)
|
||||
ipc_handles = [{"gpu-uuid": ("ipc-args",)}]
|
||||
|
||||
update_info = engine.parse_update_info(
|
||||
{
|
||||
"names": ["w1"],
|
||||
"dtype_names": ["float32"],
|
||||
"shapes": [[1]],
|
||||
"ipc_handles": ipc_handles,
|
||||
"ipc_handles_pickled": None,
|
||||
}
|
||||
)
|
||||
|
||||
assert isinstance(update_info, IPCWeightTransferUpdateInfo)
|
||||
assert update_info.ipc_handles == ipc_handles
|
||||
|
||||
def test_parse_update_info_both_handles_and_pickled_raises(self):
|
||||
"""Test that providing both ipc_handles and ipc_handles_pickled raises."""
|
||||
if torch.accelerator.device_count() < 1:
|
||||
|
||||
@@ -635,3 +635,143 @@ class TestThinkingBlockConversion:
|
||||
# Redacted thinking is ignored, normal thinking still becomes reasoning.
|
||||
assert asst.get("reasoning") == "Thinking..."
|
||||
assert asst.get("content") == "Hi!"
|
||||
|
||||
|
||||
class TestInlineSystemMessageInMessagesArray:
|
||||
"""Verify that ``role: system`` messages embedded inside the ``messages``
|
||||
array are accepted and merged with the top-level ``system`` prompt.
|
||||
|
||||
This handles clients that place system messages inside the messages array
|
||||
instead of the Anthropic-standard top-level ``system`` field.
|
||||
"""
|
||||
|
||||
def test_inline_system_merged_with_top_level_system(self):
|
||||
"""Full integration: inline system + top-level system + user message."""
|
||||
request = _make_request(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "<system-reminder>\n.....\n</system-reminder>\n\n",
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "help?",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "system",
|
||||
"content": ".....",
|
||||
},
|
||||
],
|
||||
system=[
|
||||
{
|
||||
"type": "text",
|
||||
"text": "x-anthropic-billing-header: "
|
||||
"cc_version=2.1.160.bca; cc_entrypoint=cli; cch=d1d48;",
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "You are Claude Code, Anthropic's official CLI for Claude.",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "....",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
},
|
||||
],
|
||||
tools=[],
|
||||
)
|
||||
|
||||
result = _convert(request)
|
||||
|
||||
# First message should be the merged system prompt.
|
||||
assert result.messages[0]["role"] == "system"
|
||||
# Billing header stripped, inline system appended.
|
||||
assert (
|
||||
result.messages[0]["content"]
|
||||
== "You are Claude Code, Anthropic's official CLI for Claude."
|
||||
"...."
|
||||
"....."
|
||||
)
|
||||
|
||||
# Second message should be the user message, content preserved.
|
||||
assert result.messages[1]["role"] == "user"
|
||||
user_content = result.messages[1]["content"]
|
||||
assert len(user_content) == 2
|
||||
assert user_content[0] == {
|
||||
"type": "text",
|
||||
"text": "<system-reminder>\n.....\n</system-reminder>\n\n",
|
||||
}
|
||||
assert user_content[1] == {
|
||||
"type": "text",
|
||||
"text": "help?",
|
||||
}
|
||||
|
||||
def test_inline_system_string_only(self):
|
||||
"""Only an inline system string, no top-level system."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "system", "content": "Be concise."},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
assert result.messages[0]["role"] == "system"
|
||||
assert result.messages[0]["content"] == "Be concise."
|
||||
assert result.messages[1]["role"] == "user"
|
||||
|
||||
def test_inline_system_list_content(self):
|
||||
"""Inline system with list content blocks."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Hi"},
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": "Part one. "},
|
||||
{"type": "text", "text": "Part two."},
|
||||
],
|
||||
},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
assert result.messages[0]["role"] == "system"
|
||||
assert result.messages[0]["content"] == "Part one. Part two."
|
||||
|
||||
def test_multiple_inline_system_messages(self):
|
||||
"""Multiple inline system messages should all be merged."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "system", "content": "First system."},
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "system", "content": "Second system."},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
assert result.messages[0]["role"] == "system"
|
||||
assert result.messages[0]["content"] == "First system.Second system."
|
||||
assert result.messages[1]["role"] == "user"
|
||||
|
||||
def test_inline_system_with_top_level_string(self):
|
||||
"""Top-level system is a string, inline system is also present."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "system", "content": "Inline hint."},
|
||||
],
|
||||
system="Top-level prompt.",
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
assert result.messages[0]["role"] == "system"
|
||||
assert result.messages[0]["content"] == "Top-level prompt.Inline hint."
|
||||
assert result.messages[1]["role"] == "user"
|
||||
|
||||
+8
-8
@@ -18,13 +18,13 @@ from unittest.mock import MagicMock
|
||||
import pytest
|
||||
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
|
||||
from vllm.entrypoints.openai.generative_scoring.serving import (
|
||||
from vllm.entrypoints.generate.generative_scoring.serving import (
|
||||
GenerativeScoringItemResult,
|
||||
GenerativeScoringRequest,
|
||||
GenerativeScoringResponse,
|
||||
OpenAIServingGenerativeScoring,
|
||||
ServingGenerativeScoring,
|
||||
)
|
||||
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
|
||||
from vllm.entrypoints.openai.models.protocol import BaseModelPath
|
||||
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
|
||||
from vllm.logprobs import Logprob
|
||||
@@ -86,13 +86,13 @@ def _create_mock_engine():
|
||||
return mock_engine
|
||||
|
||||
|
||||
def _create_serving(mock_engine) -> OpenAIServingGenerativeScoring:
|
||||
"""Create an OpenAIServingGenerativeScoring instance with mocks."""
|
||||
def _create_serving(mock_engine) -> ServingGenerativeScoring:
|
||||
"""Create an ServingGenerativeScoring instance with mocks."""
|
||||
models = OpenAIServingModels(
|
||||
engine_client=mock_engine,
|
||||
base_model_paths=BASE_MODEL_PATHS,
|
||||
)
|
||||
return OpenAIServingGenerativeScoring(mock_engine, models, request_logger=None)
|
||||
return ServingGenerativeScoring(mock_engine, models, request_logger=None)
|
||||
|
||||
|
||||
def _create_mock_request_output(logprobs_dict: dict[int, float]) -> RequestOutput:
|
||||
@@ -186,7 +186,7 @@ class TestProbabilityComputation:
|
||||
self, label_logprobs, apply_softmax, should_sum_to_one
|
||||
):
|
||||
"""Test probability computation for softmax and true probability modes."""
|
||||
serving = OpenAIServingGenerativeScoring.__new__(OpenAIServingGenerativeScoring)
|
||||
serving = ServingGenerativeScoring.__new__(ServingGenerativeScoring)
|
||||
probs = serving._compute_probabilities(
|
||||
label_logprobs, apply_softmax=apply_softmax
|
||||
)
|
||||
@@ -211,7 +211,7 @@ class TestProbabilityComputation:
|
||||
|
||||
def test_score_formula(self):
|
||||
"""Test the score formula: P(token[0]) / (P(token[0]) + P(token[1]))."""
|
||||
serving = OpenAIServingGenerativeScoring.__new__(OpenAIServingGenerativeScoring)
|
||||
serving = ServingGenerativeScoring.__new__(ServingGenerativeScoring)
|
||||
|
||||
# With logprobs -0.5 and -2.0, softmax gives higher prob to first token
|
||||
logprobs = {9454: -0.5, 2753: -2.0}
|
||||
+1
-1
@@ -8,7 +8,7 @@ Tests verify the full HTTP request/response flow using RemoteOpenAIServer.
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen3-0.6B"
|
||||
|
||||
@@ -14,8 +14,12 @@ from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer
|
||||
MODEL_NAME = "Qwen/Qwen2.5-Omni-3B"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
# Use module scope so the server is started once and shared across all
|
||||
# tests in this file. Starting a new vLLM server per test on XPU can
|
||||
# cause the second server startup to hang silently and exceed the
|
||||
# wait-for-server timeout, resulting in RuntimeError.
|
||||
args = [
|
||||
"--max-model-len",
|
||||
"16384",
|
||||
|
||||
@@ -6,9 +6,13 @@ from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
BatchChatCompletionRequest,
|
||||
ChatCompletionRequest,
|
||||
)
|
||||
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
|
||||
from vllm.entrypoints.openai.engine.protocol import GenerationError
|
||||
from vllm.entrypoints.openai.models.protocol import BaseModelPath
|
||||
@@ -444,3 +448,45 @@ def test_json_schema_response_format_missing_schema():
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
response_format={"type": "json_schema"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("format_value", [None, {}])
|
||||
def test_structural_tag_response_format_invalid(format_value):
|
||||
"""Malformed structural tags should be rejected during request validation."""
|
||||
with pytest.raises(
|
||||
ValidationError,
|
||||
match="Invalid response_format structural_tag",
|
||||
):
|
||||
ChatCompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
response_format={"type": "structural_tag", "format": format_value},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("format_value", [None, {}])
|
||||
def test_batch_structural_tag_response_format_invalid(format_value):
|
||||
"""Batch chat should reject malformed structural tags at request parsing."""
|
||||
with pytest.raises(
|
||||
ValidationError,
|
||||
match="Invalid response_format structural_tag",
|
||||
):
|
||||
BatchChatCompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
messages=[[{"role": "user", "content": "hello"}]],
|
||||
response_format={"type": "structural_tag", "format": format_value},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("structural_tag", ["not json", ""])
|
||||
def test_structured_outputs_structural_tag_invalid(structural_tag):
|
||||
"""Malformed direct structured_outputs structural tags should be rejected."""
|
||||
with pytest.raises(
|
||||
ValidationError,
|
||||
match="Invalid structured_outputs structural_tag",
|
||||
):
|
||||
ChatCompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
structured_outputs={"structural_tag": structural_tag},
|
||||
)
|
||||
|
||||
@@ -1935,8 +1935,10 @@ async def test_streaming_n_gt1_independent_tool_parsers():
|
||||
finished=True,
|
||||
)
|
||||
|
||||
# Collect tool-call deltas per choice from the SSE stream.
|
||||
# Collect tool-call deltas and finish_reasons per choice from the SSE
|
||||
# stream.
|
||||
tc_deltas_by_choice: dict[int, list[dict]] = {i: [] for i in range(num_choices)}
|
||||
finish_reasons_by_choice: dict[int, list[str]] = {i: [] for i in range(num_choices)}
|
||||
async for chunk_str in serving_chat.chat_completion_stream_generator(
|
||||
request=request,
|
||||
result_generator=result_generator(),
|
||||
@@ -1959,6 +1961,8 @@ async def test_streaming_n_gt1_independent_tool_parsers():
|
||||
if delta.get("tool_calls"):
|
||||
for tc in delta["tool_calls"]:
|
||||
tc_deltas_by_choice[idx].append(tc)
|
||||
if choice.get("finish_reason") is not None:
|
||||
finish_reasons_by_choice[idx].append(choice["finish_reason"])
|
||||
|
||||
# Both choices must independently produce the correct tool call.
|
||||
for choice_idx in range(num_choices):
|
||||
@@ -1984,141 +1988,11 @@ async def test_streaming_n_gt1_independent_tool_parsers():
|
||||
f"Choice {choice_idx}: expected {{'city': 'Tokyo'}}, got {parsed_args}"
|
||||
)
|
||||
|
||||
|
||||
class TestCreateRemainingArgsDelta:
|
||||
"""Tests for _create_remaining_args_delta helper function.
|
||||
|
||||
This helper is used when streaming tool calls to preserve id/type/name
|
||||
fields in the finish chunk, which would otherwise be lost.
|
||||
"""
|
||||
|
||||
def test_preserves_id_type_name(self):
|
||||
"""Test that id, type, and name are preserved from original delta."""
|
||||
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaFunctionCall,
|
||||
DeltaMessage,
|
||||
DeltaToolCall,
|
||||
reasons = finish_reasons_by_choice[choice_idx]
|
||||
assert len(reasons) == 1, (
|
||||
f"Choice {choice_idx}: expected exactly 1 finish_reason, got {reasons}"
|
||||
)
|
||||
|
||||
original_delta = DeltaMessage(
|
||||
tool_calls=[
|
||||
DeltaToolCall(
|
||||
index=0,
|
||||
id="call_abc123",
|
||||
type="function",
|
||||
function=DeltaFunctionCall(
|
||||
name="get_weather",
|
||||
arguments='{"location": "Paris"}',
|
||||
),
|
||||
)
|
||||
]
|
||||
assert reasons[0] == "tool_calls", (
|
||||
f"Choice {choice_idx}: expected finish_reason='tool_calls', "
|
||||
f"got '{reasons[0]}'"
|
||||
)
|
||||
|
||||
result = OpenAIServingChat._create_remaining_args_delta(
|
||||
original_delta, '", "unit": "celsius"}', 0
|
||||
)
|
||||
|
||||
assert len(result.tool_calls) == 1
|
||||
tc = result.tool_calls[0]
|
||||
assert tc.index == 0
|
||||
assert tc.id == "call_abc123"
|
||||
assert tc.type == "function"
|
||||
assert tc.function.name == "get_weather"
|
||||
assert tc.function.arguments == '", "unit": "celsius"}'
|
||||
|
||||
def test_matches_by_index(self):
|
||||
"""Test that the correct tool call is matched by index."""
|
||||
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaFunctionCall,
|
||||
DeltaMessage,
|
||||
DeltaToolCall,
|
||||
)
|
||||
|
||||
original_delta = DeltaMessage(
|
||||
tool_calls=[
|
||||
DeltaToolCall(
|
||||
index=0,
|
||||
id="call_first",
|
||||
type="function",
|
||||
function=DeltaFunctionCall(name="func_a", arguments="{}"),
|
||||
),
|
||||
DeltaToolCall(
|
||||
index=1,
|
||||
id="call_second",
|
||||
type="function",
|
||||
function=DeltaFunctionCall(name="func_b", arguments="{}"),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
result = OpenAIServingChat._create_remaining_args_delta(
|
||||
original_delta, '{"extra": true}', 1
|
||||
)
|
||||
|
||||
assert len(result.tool_calls) == 1
|
||||
tc = result.tool_calls[0]
|
||||
assert tc.index == 1
|
||||
assert tc.id == "call_second"
|
||||
assert tc.function.name == "func_b"
|
||||
|
||||
def test_no_matching_tool_call(self):
|
||||
"""Test graceful handling when no matching tool call is found."""
|
||||
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaFunctionCall,
|
||||
DeltaMessage,
|
||||
DeltaToolCall,
|
||||
)
|
||||
|
||||
original_delta = DeltaMessage(
|
||||
tool_calls=[
|
||||
DeltaToolCall(
|
||||
index=0,
|
||||
id="call_zero",
|
||||
type="function",
|
||||
function=DeltaFunctionCall(name="func", arguments="{}"),
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
result = OpenAIServingChat._create_remaining_args_delta(
|
||||
original_delta, '{"arg": 1}', 5
|
||||
)
|
||||
|
||||
assert len(result.tool_calls) == 1
|
||||
tc = result.tool_calls[0]
|
||||
assert tc.index == 5
|
||||
assert tc.id is None
|
||||
assert tc.type is None
|
||||
assert tc.function.name is None
|
||||
assert tc.function.arguments == '{"arg": 1}'
|
||||
|
||||
def test_function_is_none(self):
|
||||
"""Test handling when original tool call has no function."""
|
||||
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
|
||||
from vllm.entrypoints.openai.engine.protocol import DeltaMessage, DeltaToolCall
|
||||
|
||||
original_delta = DeltaMessage(
|
||||
tool_calls=[
|
||||
DeltaToolCall(
|
||||
index=0,
|
||||
id="call_nofunc",
|
||||
type="function",
|
||||
function=None,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
result = OpenAIServingChat._create_remaining_args_delta(
|
||||
original_delta, '{"data": "value"}', 0
|
||||
)
|
||||
|
||||
assert len(result.tool_calls) == 1
|
||||
tc = result.tool_calls[0]
|
||||
assert tc.index == 0
|
||||
assert tc.id == "call_nofunc"
|
||||
assert tc.type == "function"
|
||||
assert tc.function.name is None
|
||||
assert tc.function.arguments == '{"data": "value"}'
|
||||
|
||||
@@ -6,6 +6,7 @@ from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.entrypoints.openai.completion.protocol import CompletionRequest
|
||||
@@ -302,6 +303,36 @@ def test_json_schema_response_format_missing_schema():
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("format_value", [None, {}])
|
||||
def test_structural_tag_response_format_invalid(format_value):
|
||||
"""Malformed structural tags should be rejected during request validation."""
|
||||
with pytest.raises(
|
||||
ValidationError,
|
||||
match="Invalid response_format structural_tag",
|
||||
):
|
||||
CompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
prompt="Test prompt",
|
||||
max_tokens=10,
|
||||
response_format={"type": "structural_tag", "format": format_value},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("structural_tag", ["not json", ""])
|
||||
def test_structured_outputs_structural_tag_invalid(structural_tag):
|
||||
"""Malformed direct structured_outputs structural tags should be rejected."""
|
||||
with pytest.raises(
|
||||
ValidationError,
|
||||
match="Invalid structured_outputs structural_tag",
|
||||
):
|
||||
CompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
prompt="Test prompt",
|
||||
max_tokens=10,
|
||||
structured_outputs={"structural_tag": structural_tag},
|
||||
)
|
||||
|
||||
|
||||
def test_negative_prompt_token_ids_nested():
|
||||
"""Negative token IDs in prompt (nested list) should raise validation error."""
|
||||
with pytest.raises(Exception, match="greater than or equal to 0"):
|
||||
|
||||
@@ -0,0 +1,382 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for ResponsesParser with the unified Parser interface.
|
||||
|
||||
These tests verify that ResponsesParser correctly delegates to the unified
|
||||
Parser (via extract_response_outputs) instead of calling separate
|
||||
ReasoningParser / ToolParser instances directly.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaMessage,
|
||||
ExtractedToolCallInformation,
|
||||
FunctionCall,
|
||||
ToolCall,
|
||||
)
|
||||
from vllm.entrypoints.openai.parser.responses_parser import (
|
||||
ResponsesParser,
|
||||
get_responses_parser_for_simple_context,
|
||||
)
|
||||
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
|
||||
from vllm.outputs import CompletionOutput
|
||||
from vllm.parser.abstract_parser import DelegatingParser
|
||||
|
||||
pytestmark = pytest.mark.skip_global_cleanup
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test parser stubs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _NoOpParser(DelegatingParser):
|
||||
"""Parser that extracts no reasoning and no tool calls."""
|
||||
|
||||
def is_reasoning_end(self, input_ids: list[int]) -> bool:
|
||||
return False
|
||||
|
||||
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
|
||||
return input_ids
|
||||
|
||||
def extract_reasoning(self, model_output, request):
|
||||
return None, model_output
|
||||
|
||||
def extract_reasoning_streaming(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
def extract_tool_calls(self, model_output, request):
|
||||
return ExtractedToolCallInformation(
|
||||
tools_called=False, tool_calls=[], content=model_output
|
||||
)
|
||||
|
||||
def extract_tool_calls_streaming(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
def parse_delta(self, *args, **kwargs) -> DeltaMessage | None:
|
||||
return None
|
||||
|
||||
|
||||
class _ReasoningOnlyParser(DelegatingParser):
|
||||
"""Parser that extracts reasoning but no tool calls."""
|
||||
|
||||
def is_reasoning_end(self, input_ids: list[int]) -> bool:
|
||||
return False
|
||||
|
||||
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
|
||||
return input_ids
|
||||
|
||||
def extract_reasoning(self, model_output, request):
|
||||
if "<think>" in model_output and "</think>" in model_output:
|
||||
start = model_output.index("<think>") + len("<think>")
|
||||
end = model_output.index("</think>")
|
||||
reasoning = model_output[start:end]
|
||||
content = model_output[end + len("</think>") :]
|
||||
return reasoning, content.strip() or None
|
||||
return None, model_output
|
||||
|
||||
def extract_reasoning_streaming(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
def extract_tool_calls(self, model_output, request):
|
||||
return ExtractedToolCallInformation(
|
||||
tools_called=False, tool_calls=[], content=model_output
|
||||
)
|
||||
|
||||
def extract_tool_calls_streaming(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
def parse_delta(self, *args, **kwargs) -> DeltaMessage | None:
|
||||
return None
|
||||
|
||||
|
||||
class _StubToolParser:
|
||||
"""Minimal tool parser stub that always returns a hardcoded tool call."""
|
||||
|
||||
supports_required_and_named = False
|
||||
|
||||
def __init__(self, tokenizer=None, tools=None):
|
||||
pass
|
||||
|
||||
def extract_tool_calls(self, model_output, request):
|
||||
return ExtractedToolCallInformation(
|
||||
tools_called=True,
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
id="call_123",
|
||||
type="function",
|
||||
function=FunctionCall(
|
||||
name="get_weather",
|
||||
arguments='{"location": "Paris"}',
|
||||
),
|
||||
)
|
||||
],
|
||||
content=None,
|
||||
)
|
||||
|
||||
def extract_tool_calls_streaming(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
def adjust_request(self, request):
|
||||
return request
|
||||
|
||||
|
||||
class _ToolCallingParser(DelegatingParser):
|
||||
"""Parser that extracts a hardcoded tool call from any input."""
|
||||
|
||||
def __init__(self, tokenizer, *args, **kwargs):
|
||||
super().__init__(tokenizer)
|
||||
self._tool_parser = _StubToolParser()
|
||||
|
||||
def is_reasoning_end(self, input_ids: list[int]) -> bool:
|
||||
return False
|
||||
|
||||
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
|
||||
return input_ids
|
||||
|
||||
def extract_reasoning(self, model_output, request):
|
||||
return None, model_output
|
||||
|
||||
def extract_reasoning_streaming(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
def extract_tool_calls_streaming(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
def parse_delta(self, *args, **kwargs) -> DeltaMessage | None:
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_request(**overrides) -> ResponsesRequest:
|
||||
defaults = {"model": "test-model", "input": "test"}
|
||||
defaults.update(overrides)
|
||||
return ResponsesRequest.model_validate(defaults)
|
||||
|
||||
|
||||
def _make_output(
|
||||
text: str = "Hello, world!",
|
||||
token_ids: Sequence[int] = (1, 2, 3),
|
||||
finish_reason: str = "stop",
|
||||
) -> CompletionOutput:
|
||||
return CompletionOutput(
|
||||
index=0,
|
||||
text=text,
|
||||
token_ids=list(token_ids),
|
||||
cumulative_logprob=None,
|
||||
logprobs=None,
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
|
||||
|
||||
def _make_parser(parser_cls, **overrides):
|
||||
defaults = dict(
|
||||
tokenizer=MagicMock(),
|
||||
parser_cls=parser_cls,
|
||||
response_messages=[],
|
||||
request=_make_request(),
|
||||
chat_template=None,
|
||||
chat_template_content_format="auto",
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return ResponsesParser(**defaults)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: basic text passthrough
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_process_text_with_parser():
|
||||
"""Parser with no reasoning/tools returns a single message item."""
|
||||
parser = _make_parser(_NoOpParser)
|
||||
parser.process(_make_output(text="Hello!"))
|
||||
|
||||
assert len(parser.response_messages) == 1
|
||||
msg = parser.response_messages[0]
|
||||
assert msg.type == "message"
|
||||
assert msg.content[0].text == "Hello!"
|
||||
|
||||
|
||||
def test_process_text_without_parser():
|
||||
"""parser_cls=None falls back to plain text wrapping."""
|
||||
parser = _make_parser(None)
|
||||
parser.process(_make_output(text="Hello!"))
|
||||
|
||||
assert len(parser.response_messages) == 1
|
||||
msg = parser.response_messages[0]
|
||||
assert msg.type == "message"
|
||||
assert msg.content[0].text == "Hello!"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: empty / whitespace output
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_process_empty_text_without_parser():
|
||||
"""Empty text with no parser produces no output items."""
|
||||
parser = _make_parser(None)
|
||||
parser.process(_make_output(text=""))
|
||||
|
||||
assert len(parser.response_messages) == 0
|
||||
|
||||
|
||||
def test_process_empty_text_with_parser():
|
||||
"""Empty text with parser produces no output items."""
|
||||
parser = _make_parser(_NoOpParser)
|
||||
parser.process(_make_output(text=""))
|
||||
|
||||
assert len(parser.response_messages) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: reasoning extraction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_process_extracts_reasoning():
|
||||
"""Parser that finds reasoning produces both reasoning and message items."""
|
||||
parser = _make_parser(_ReasoningOnlyParser)
|
||||
parser.process(_make_output(text="<think>Let me check</think>The answer is 42"))
|
||||
|
||||
types = [m.type for m in parser.response_messages]
|
||||
assert "reasoning" in types
|
||||
assert "message" in types
|
||||
|
||||
reasoning_item = next(m for m in parser.response_messages if m.type == "reasoning")
|
||||
assert reasoning_item.content[0].text == "Let me check"
|
||||
|
||||
message_item = next(m for m in parser.response_messages if m.type == "message")
|
||||
assert message_item.content[0].text == "The answer is 42"
|
||||
|
||||
|
||||
def test_process_reasoning_only_no_content():
|
||||
"""When reasoning consumes all text, only a reasoning item is produced."""
|
||||
parser = _make_parser(_ReasoningOnlyParser)
|
||||
parser.process(_make_output(text="<think>Just thinking</think>"))
|
||||
|
||||
types = [m.type for m in parser.response_messages]
|
||||
assert "reasoning" in types
|
||||
assert "message" not in types
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: tool call extraction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_process_extracts_tool_calls():
|
||||
"""Parser that finds tool calls produces function_call items."""
|
||||
request = _make_request(
|
||||
tool_choice="auto",
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_weather",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
}
|
||||
],
|
||||
)
|
||||
parser = _make_parser(_ToolCallingParser, request=request, enable_auto_tools=True)
|
||||
parser.process(_make_output(text="calling tool"))
|
||||
|
||||
types = [m.type for m in parser.response_messages]
|
||||
assert "function_call" in types
|
||||
|
||||
tool_item = next(m for m in parser.response_messages if m.type == "function_call")
|
||||
assert tool_item.name == "get_weather"
|
||||
assert tool_item.arguments == '{"location": "Paris"}'
|
||||
assert tool_item.status == "completed"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: finish_reason tracking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_finish_reason_tracked():
|
||||
"""finish_reason from CompletionOutput is stored on the parser."""
|
||||
parser = _make_parser(_NoOpParser)
|
||||
assert parser.finish_reason is None
|
||||
|
||||
parser.process(_make_output(finish_reason="stop"))
|
||||
assert parser.finish_reason == "stop"
|
||||
|
||||
parser.process(_make_output(finish_reason="length"))
|
||||
assert parser.finish_reason == "length"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: multi-turn accumulation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_multi_turn_accumulation():
|
||||
"""Multiple process() calls accumulate response_messages."""
|
||||
parser = _make_parser(_NoOpParser)
|
||||
|
||||
parser.process(_make_output(text="First turn"))
|
||||
parser.process(_make_output(text="Second turn"))
|
||||
|
||||
assert len(parser.response_messages) == 2
|
||||
texts = [m.content[0].text for m in parser.response_messages]
|
||||
assert texts == ["First turn", "Second turn"]
|
||||
|
||||
|
||||
def test_num_init_messages_offset():
|
||||
"""Initial messages are preserved and offset works correctly."""
|
||||
init_messages = [MagicMock(type="message")]
|
||||
parser = _make_parser(_NoOpParser, response_messages=init_messages)
|
||||
|
||||
assert parser.num_init_messages == 1
|
||||
|
||||
parser.process(_make_output(text="New output"))
|
||||
|
||||
assert len(parser.response_messages) == 2
|
||||
items = parser.make_response_output_items_from_parsable_context()
|
||||
assert len(items) == 1
|
||||
assert items[0].type == "message"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: factory function
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_factory_function_creates_parser():
|
||||
"""get_responses_parser_for_simple_context returns a working parser."""
|
||||
rp = get_responses_parser_for_simple_context(
|
||||
tokenizer=MagicMock(),
|
||||
parser_cls=_NoOpParser,
|
||||
response_messages=[],
|
||||
request=_make_request(),
|
||||
chat_template=None,
|
||||
chat_template_content_format="auto",
|
||||
)
|
||||
assert isinstance(rp, ResponsesParser)
|
||||
|
||||
rp.process(_make_output(text="Works!"))
|
||||
assert len(rp.response_messages) == 1
|
||||
|
||||
|
||||
def test_factory_function_none_parser():
|
||||
"""Factory function works with parser_cls=None."""
|
||||
rp = get_responses_parser_for_simple_context(
|
||||
tokenizer=MagicMock(),
|
||||
parser_cls=None,
|
||||
response_messages=[],
|
||||
request=_make_request(),
|
||||
chat_template=None,
|
||||
chat_template_content_format="auto",
|
||||
)
|
||||
assert isinstance(rp, ResponsesParser)
|
||||
assert rp.parser_instance is None
|
||||
@@ -4,7 +4,6 @@
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.entrypoints.openai.engine.serving import OpenAIServing
|
||||
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
|
||||
from vllm.parser.abstract_parser import DelegatingParser
|
||||
|
||||
@@ -32,11 +31,8 @@ class _DummyDelegatingParser(DelegatingParser):
|
||||
):
|
||||
return None
|
||||
|
||||
def extract_tool_calls(self, model_output: str, request):
|
||||
return None
|
||||
|
||||
|
||||
def test_parse_tool_calls_from_content_allows_named_tool_choice_with_none_content():
|
||||
def test_chat_completion_named_tool_choice_with_none_content():
|
||||
request = ChatCompletionRequest.model_validate(
|
||||
{
|
||||
"model": "test-model",
|
||||
@@ -53,17 +49,15 @@ def test_parse_tool_calls_from_content_allows_named_tool_choice_with_none_conten
|
||||
"tool_choice": {"type": "function", "function": {"name": "get_weather"}},
|
||||
}
|
||||
)
|
||||
parser = _DummyDelegatingParser(tokenizer=None)
|
||||
|
||||
tool_calls, content = OpenAIServing._parse_tool_calls_from_content(
|
||||
request=request,
|
||||
tokenizer=None,
|
||||
enable_auto_tools=True,
|
||||
tool_parser_cls=None,
|
||||
tool_calls, content = parser._extract_tool_calls(
|
||||
content=None,
|
||||
request=request,
|
||||
enable_auto_tools=True,
|
||||
)
|
||||
|
||||
assert content is None
|
||||
assert tool_calls is not None
|
||||
assert tool_calls == []
|
||||
|
||||
|
||||
|
||||
+1
-1
@@ -37,7 +37,7 @@ def server():
|
||||
"--max-num-seqs",
|
||||
"128",
|
||||
"--worker-extension-cls",
|
||||
"tests.entrypoints.rpc.test_collective_rpc.TestWorkerExtension",
|
||||
"tests.entrypoints.serve.dev.rpc.test_collective_rpc.TestWorkerExtension",
|
||||
]
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME,
|
||||
@@ -48,6 +48,7 @@ class MockUpdateInfo(WeightTransferUpdateInfo):
|
||||
names: list[str] | None = None
|
||||
dtype_names: list[str] | None = None
|
||||
shapes: list[list[int]] | None = None
|
||||
num_updates_list: list[int] | None = None
|
||||
|
||||
|
||||
class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]):
|
||||
@@ -87,6 +88,15 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
|
||||
# (In real implementation, this would receive and load actual weights)
|
||||
load_weights([])
|
||||
|
||||
def receive_sparse_weights(
|
||||
self,
|
||||
update_info: MockUpdateInfo,
|
||||
apply_patches: Callable[[list], None],
|
||||
) -> None:
|
||||
MockWeightTransferEngine.receive_weights_called = True
|
||||
MockWeightTransferEngine.last_update_info = update_info
|
||||
apply_patches([])
|
||||
|
||||
def shutdown(self) -> None:
|
||||
MockWeightTransferEngine.shutdown_called = True
|
||||
|
||||
@@ -198,8 +208,6 @@ def test_update_weights_calls_engine():
|
||||
llm.init_weight_transfer_engine(
|
||||
WeightTransferInitRequest(init_info={"test_param": "init"})
|
||||
)
|
||||
|
||||
# Start weight update (required before update_weights)
|
||||
llm.start_weight_update(is_checkpoint_format=True)
|
||||
|
||||
# Call update_weights
|
||||
@@ -232,14 +240,67 @@ def test_update_weights_calls_engine():
|
||||
assert dtypes == test_dtypes
|
||||
assert shapes == test_shapes
|
||||
|
||||
# Finish weight update
|
||||
llm.finish_weight_update()
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_update_weights_passes_sparse_metadata():
|
||||
"""Test sparse update metadata is forwarded unchanged to the engine."""
|
||||
if torch.accelerator.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU for this test")
|
||||
|
||||
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
|
||||
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
|
||||
|
||||
with patch(
|
||||
"vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
|
||||
mock_create_engine,
|
||||
):
|
||||
llm = LLM(
|
||||
model=MODEL_NAME,
|
||||
enforce_eager=True,
|
||||
load_format="dummy",
|
||||
tensor_parallel_size=1,
|
||||
weight_transfer_config=WeightTransferConfig(backend="nccl"),
|
||||
)
|
||||
|
||||
llm.init_weight_transfer_engine(
|
||||
WeightTransferInitRequest(init_info={"test_param": "init"})
|
||||
)
|
||||
llm.start_weight_update(is_checkpoint_format=False)
|
||||
|
||||
llm.update_weights(
|
||||
WeightTransferUpdateRequest(
|
||||
update_info={
|
||||
"names": ["layer.weight"],
|
||||
"dtype_names": ["bfloat16"],
|
||||
"shapes": [[100]],
|
||||
"num_updates_list": [3],
|
||||
"update_kind": "sparse_flat",
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
def check_sparse_update_called(self):
|
||||
engine = self.weight_transfer_engine
|
||||
if not engine.receive_weights_called:
|
||||
return None
|
||||
info = engine.last_update_info
|
||||
return (
|
||||
info.update_kind,
|
||||
info.num_updates_list,
|
||||
)
|
||||
|
||||
results = llm.collective_rpc(check_sparse_update_called)
|
||||
for result in results:
|
||||
assert result == ("sparse_flat", [3])
|
||||
|
||||
llm.finish_weight_update()
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_full_weight_transfer_flow():
|
||||
"""Test the complete weight transfer flow:
|
||||
init -> start -> update -> finish."""
|
||||
"""Test the complete weight transfer flow: init -> start -> update -> finish."""
|
||||
if torch.accelerator.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU for this test")
|
||||
|
||||
|
||||
@@ -0,0 +1,314 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import functools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not current_platform.is_cpu():
|
||||
pytest.skip("skipping CPU-only tests", allow_module_level=True)
|
||||
|
||||
set_random_seed(12345)
|
||||
|
||||
NUM_HEADS = [
|
||||
(2, 4),
|
||||
(4, 4),
|
||||
]
|
||||
HEAD_DIMS = [
|
||||
(32, 32),
|
||||
(64, 32),
|
||||
]
|
||||
CHUNK_SIZE = 64
|
||||
PREFILL_SEQ_LENS = [
|
||||
[1],
|
||||
[1, 2, 3],
|
||||
[CHUNK_SIZE - 1],
|
||||
[CHUNK_SIZE],
|
||||
[CHUNK_SIZE + 1],
|
||||
[CHUNK_SIZE - 1, CHUNK_SIZE, CHUNK_SIZE + 1],
|
||||
[2 * CHUNK_SIZE - 1, 2 * CHUNK_SIZE, 2 * CHUNK_SIZE + 1],
|
||||
[4 * CHUNK_SIZE + 17],
|
||||
]
|
||||
DECODE_BATCH_SIZES = [1, 3, 5]
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=128, typed=False)
|
||||
def tensor_cache(
|
||||
elem_num: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
tensor = torch.rand(elem_num, dtype=dtype)
|
||||
return tensor
|
||||
|
||||
|
||||
def ref_l2norm(
|
||||
x: torch.Tensor,
|
||||
dim: int = -1,
|
||||
eps: float = 1e-5,
|
||||
) -> torch.Tensor:
|
||||
inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
|
||||
return x * inv_norm
|
||||
|
||||
|
||||
def ref_gdn_gating(
|
||||
A_log: torch.Tensor,
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
dt_bias: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
softplus_x = F.softplus(a.float() + dt_bias.float(), beta=1.0, threshold=20.0)
|
||||
g = -torch.exp(A_log.float()) * softplus_x
|
||||
beta = torch.sigmoid(b.float()).to(dtype=b.dtype)
|
||||
return g, beta
|
||||
|
||||
|
||||
def ref_gated_delta_rule(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
A_log: torch.Tensor,
|
||||
dt_bias: torch.Tensor,
|
||||
initial_state: torch.Tensor,
|
||||
cu_seqlens: torch.Tensor,
|
||||
use_qk_l2norm_in_kernel: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
|
||||
out = torch.empty_like(value)
|
||||
final_state = torch.empty_like(initial_state)
|
||||
|
||||
for seq_idx in range(cu_seqlens.numel() - 1):
|
||||
begin = int(cu_seqlens[seq_idx].item())
|
||||
end = int(cu_seqlens[seq_idx + 1].item())
|
||||
q_seq = query[:, begin:end]
|
||||
k_seq = key[:, begin:end]
|
||||
v_seq = value[:, begin:end]
|
||||
g_seq = g[begin:end].unsqueeze(0)
|
||||
beta_seq = beta[begin:end].unsqueeze(0)
|
||||
initial_dtype = q_seq.dtype
|
||||
|
||||
if use_qk_l2norm_in_kernel:
|
||||
q_seq = ref_l2norm(q_seq, dim=-1)
|
||||
k_seq = ref_l2norm(k_seq, dim=-1)
|
||||
|
||||
if q_seq.shape[2] != v_seq.shape[2]:
|
||||
repeat_factor = v_seq.shape[2] // q_seq.shape[2]
|
||||
q_seq = q_seq.repeat_interleave(repeat_factor, dim=2)
|
||||
k_seq = k_seq.repeat_interleave(repeat_factor, dim=2)
|
||||
|
||||
q_seq, k_seq, v_seq, beta_seq, g_seq = [
|
||||
x.transpose(1, 2).contiguous().to(torch.float32)
|
||||
for x in (q_seq, k_seq, v_seq, beta_seq, g_seq)
|
||||
]
|
||||
|
||||
batch_size, num_heads, seq_len, head_dim = q_seq.shape
|
||||
v_head_dim = v_seq.shape[-1]
|
||||
q_seq = q_seq * (1 / (head_dim**0.5))
|
||||
out_seq = torch.empty(
|
||||
batch_size,
|
||||
num_heads,
|
||||
seq_len,
|
||||
v_head_dim,
|
||||
dtype=v_seq.dtype,
|
||||
)
|
||||
state = initial_state[seq_idx : seq_idx + 1].to(v_seq)
|
||||
|
||||
for token_idx in range(seq_len):
|
||||
q_t = q_seq[:, :, token_idx]
|
||||
k_t = k_seq[:, :, token_idx]
|
||||
v_t = v_seq[:, :, token_idx]
|
||||
g_t = g_seq[:, :, token_idx].exp().unsqueeze(-1).unsqueeze(-1)
|
||||
beta_t = beta_seq[:, :, token_idx].unsqueeze(-1)
|
||||
|
||||
state = state * g_t
|
||||
kv_mem = (state * k_t.unsqueeze(-2)).sum(dim=-1)
|
||||
delta = (v_t - kv_mem) * beta_t
|
||||
state = state + delta.unsqueeze(-1) * k_t.unsqueeze(-2)
|
||||
out_seq[:, :, token_idx] = (state * q_t.unsqueeze(-2)).sum(dim=-1)
|
||||
|
||||
out[:, begin:end] = out_seq.transpose(1, 2).contiguous().to(initial_dtype)
|
||||
final_state[seq_idx] = state.squeeze(0)
|
||||
|
||||
return out, final_state
|
||||
|
||||
|
||||
def gdn_inputs(
|
||||
num_tokens: int,
|
||||
num_heads: tuple[int, int],
|
||||
head_dims: tuple[int, int],
|
||||
) -> tuple[torch.Tensor, ...]:
|
||||
num_qk_heads, num_v_heads = num_heads
|
||||
head_dim, v_head_dim = head_dims
|
||||
q_shape = (1, num_tokens, num_qk_heads, head_dim)
|
||||
q_numel = num_tokens * num_qk_heads * head_dim
|
||||
q = tensor_cache(q_numel, torch.bfloat16).view(q_shape)
|
||||
k = tensor_cache(q_numel, torch.bfloat16).view(q_shape)
|
||||
|
||||
v_shape = (1, num_tokens, num_v_heads, v_head_dim)
|
||||
v = tensor_cache(num_tokens * num_v_heads * v_head_dim, torch.bfloat16).view(
|
||||
v_shape
|
||||
)
|
||||
|
||||
gate_shape = (num_tokens, num_v_heads)
|
||||
gate_numel = num_tokens * num_v_heads
|
||||
a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
|
||||
b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
|
||||
A_log = tensor_cache(num_v_heads, torch.float32)
|
||||
dt_bias = tensor_cache(num_v_heads, torch.bfloat16)
|
||||
return q, k, v, a, b, A_log, dt_bias
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 9])
|
||||
@pytest.mark.parametrize("num_v_heads", [4, 8])
|
||||
@torch.inference_mode()
|
||||
def test_fused_gdn_gating_cpu(
|
||||
num_tokens: int,
|
||||
num_v_heads: int,
|
||||
) -> None:
|
||||
gate_shape = (num_tokens, num_v_heads)
|
||||
gate_numel = num_tokens * num_v_heads
|
||||
a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
|
||||
b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
|
||||
A_log = tensor_cache(num_v_heads, torch.float32)
|
||||
dt_bias = tensor_cache(num_v_heads, torch.bfloat16)
|
||||
|
||||
g_ref, beta_ref = ref_gdn_gating(A_log, a, b, dt_bias)
|
||||
g, beta = ops.fused_gdn_gating_cpu(A_log, a, b, dt_bias)
|
||||
|
||||
torch.testing.assert_close(g, g_ref.unsqueeze(0), atol=1e-4, rtol=1e-4)
|
||||
torch.testing.assert_close(
|
||||
beta.float(), beta_ref.unsqueeze(0).float(), atol=5e-3, rtol=5e-3
|
||||
)
|
||||
|
||||
|
||||
# decode path
|
||||
@pytest.mark.parametrize("batch_size", DECODE_BATCH_SIZES)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_dims", HEAD_DIMS)
|
||||
@torch.inference_mode()
|
||||
def test_fused_sigmoid_gating_delta_rule_update_cpu(
|
||||
batch_size: int,
|
||||
num_heads: tuple[int, int],
|
||||
head_dims: tuple[int, int],
|
||||
) -> None:
|
||||
q, k, v, a, b, A_log, dt_bias = gdn_inputs(
|
||||
num_tokens=batch_size,
|
||||
num_heads=num_heads,
|
||||
head_dims=head_dims,
|
||||
)
|
||||
_, num_v_heads = num_heads
|
||||
head_dim, v_head_dim = head_dims
|
||||
state_indices = torch.arange(batch_size, dtype=torch.int32)
|
||||
cu_seqlens = torch.arange(batch_size + 1, dtype=torch.int32)
|
||||
state_shape = (batch_size, num_v_heads, head_dim, v_head_dim)
|
||||
state = tensor_cache(
|
||||
batch_size * num_v_heads * head_dim * v_head_dim, torch.float32
|
||||
).view(state_shape)
|
||||
state_ref = state[state_indices].transpose(-1, -2).contiguous()
|
||||
|
||||
out_ref, final_state_ref = ref_gated_delta_rule(
|
||||
query=q,
|
||||
key=k,
|
||||
value=v,
|
||||
a=a,
|
||||
b=b,
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
initial_state=state_ref,
|
||||
cu_seqlens=cu_seqlens,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
out_ref = out_ref.transpose(0, 1).contiguous()
|
||||
|
||||
state_out = state.clone()
|
||||
out = ops.fused_sigmoid_gating_delta_rule_update_cpu(
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
a=a,
|
||||
b=b,
|
||||
initial_state_source=state_out,
|
||||
initial_state_indices=state_indices,
|
||||
cu_seqlens=cu_seqlens,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(
|
||||
state_out[state_indices].transpose(-1, -2),
|
||||
final_state_ref,
|
||||
atol=1e-2,
|
||||
rtol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
# prefill path
|
||||
@pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_dims", HEAD_DIMS)
|
||||
@torch.inference_mode()
|
||||
def test_chunk_gated_delta_rule_cpu(
|
||||
seq_lens: list[int],
|
||||
num_heads: tuple[int, int],
|
||||
head_dims: tuple[int, int],
|
||||
) -> None:
|
||||
total_tokens = sum(seq_lens)
|
||||
q, k, v, a, b, A_log, dt_bias = gdn_inputs(
|
||||
num_tokens=total_tokens,
|
||||
num_heads=num_heads,
|
||||
head_dims=head_dims,
|
||||
)
|
||||
_, num_v_heads = num_heads
|
||||
head_dim, v_head_dim = head_dims
|
||||
cu_seqlens = torch.tensor(
|
||||
[0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32
|
||||
)
|
||||
initial_state_shape = (len(seq_lens), num_v_heads, head_dim, v_head_dim)
|
||||
initial_state = tensor_cache(
|
||||
len(seq_lens) * num_v_heads * head_dim * v_head_dim, torch.float32
|
||||
).view(initial_state_shape)
|
||||
initial_state_ref = initial_state.transpose(-1, -2).contiguous()
|
||||
|
||||
out_ref, final_state_ref = ref_gated_delta_rule(
|
||||
query=q,
|
||||
key=k,
|
||||
value=v,
|
||||
a=a,
|
||||
b=b,
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
initial_state=initial_state_ref,
|
||||
cu_seqlens=cu_seqlens,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
|
||||
out, final_state = ops.chunk_gated_delta_rule_cpu(
|
||||
query=q,
|
||||
key=k,
|
||||
value=v,
|
||||
g=g.unsqueeze(0),
|
||||
beta=beta.unsqueeze(0),
|
||||
initial_state=initial_state,
|
||||
output_final_state=True,
|
||||
cu_seqlens=cu_seqlens,
|
||||
head_first=False,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(
|
||||
final_state.transpose(-1, -2),
|
||||
final_state_ref,
|
||||
atol=1e-2,
|
||||
rtol=1e-2,
|
||||
)
|
||||
@@ -1,10 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from importlib.metadata import version
|
||||
|
||||
import pytest
|
||||
from packaging.version import Version
|
||||
|
||||
import vllm
|
||||
from vllm.assets.image import ImageAsset
|
||||
@@ -13,14 +10,6 @@ from vllm.platforms import current_platform
|
||||
|
||||
from ..utils import multi_gpu_test
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
Version("5.0") <= Version(version("transformers")),
|
||||
reason=(
|
||||
"MiniCPMV custom processor uses tokenizer.im_start_id which is not "
|
||||
"available on TokenizersBackend in transformers v5.0+"
|
||||
),
|
||||
)
|
||||
|
||||
MODEL_PATH = "openbmb/MiniCPM-Llama3-V-2_5"
|
||||
|
||||
PROMPT_TEMPLATE = (
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user