forked from Karylab-cklius/vllm
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Wentao YeandGitHub
|
bfb22a7d1d | ||
|
yewentao256
|
8ea74c05c8 |
@@ -17,26 +17,6 @@ steps:
|
||||
--target test
|
||||
--no-cache
|
||||
--progress plain .
|
||||
- |
|
||||
docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
|
||||
if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
|
||||
if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
|
||||
if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
|
||||
if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
|
||||
command -v python3
|
||||
command -v uv
|
||||
command -v pytest
|
||||
if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
|
||||
echo No ROCm CLI found in image >&2
|
||||
exit 1
|
||||
fi
|
||||
python3 - <<PY
|
||||
import torch, vllm
|
||||
print(torch.__version__)
|
||||
print(vllm.__version__)
|
||||
PY
|
||||
echo AMD image smoke OK
|
||||
'
|
||||
- docker push "rocm/vllm-ci:${BUILDKITE_COMMIT}"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -64,7 +64,7 @@ steps:
|
||||
- vllm/v1/worker/gpu/
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
|
||||
uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
|
||||
VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model"
|
||||
|
||||
@@ -74,10 +74,11 @@ steps:
|
||||
no_plugin: true
|
||||
source_file_dependencies:
|
||||
- csrc/cpu/
|
||||
- vllm/model_executor/layers/quantization/cpu_wna16.py
|
||||
- vllm/model_executor/layers/quantization/auto_gptq.py
|
||||
- vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
|
||||
- vllm/model_executor/kernels/linear/mixed_precision/cpu.py
|
||||
- vllm/model_executor/kernels/linear/scaled_mm/cpu.py
|
||||
- vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
|
||||
- vllm/model_executor/layers/quantization/kernels/mixed_precision/cpu.py
|
||||
- vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
|
||||
- tests/quantization/test_compressed_tensors.py
|
||||
- tests/quantization/test_cpu_wna16.py
|
||||
|
||||
@@ -98,21 +98,3 @@ steps:
|
||||
limit: 2
|
||||
- exit_status: -10 # Agent was lost
|
||||
limit: 2
|
||||
|
||||
- label: ":docker: Build arm64 image"
|
||||
key: arm64-image-build
|
||||
depends_on: []
|
||||
source_file_dependencies:
|
||||
- ".buildkite/image_build/image_build.yaml"
|
||||
- ".buildkite/image_build/image_build_arm64.sh"
|
||||
- "docker/Dockerfile"
|
||||
commands:
|
||||
- .buildkite/image_build/image_build_arm64.sh $REGISTRY $REPO $BUILDKITE_COMMIT
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: -1 # Agent was lost
|
||||
limit: 2
|
||||
- exit_status: -10 # Agent was lost
|
||||
limit: 2
|
||||
|
||||
@@ -1,37 +0,0 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
if [[ $# -lt 3 ]]; then
|
||||
echo "Usage: $0 <registry> <repo> <commit>"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
REGISTRY=$1
|
||||
REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
|
||||
# authenticate with AWS ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
|
||||
|
||||
# skip build if image already exists
|
||||
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64) ]]; then
|
||||
echo "Image not found, proceeding with build..."
|
||||
else
|
||||
echo "Image found"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# build (Grace/GH200 is the arm64 GPU target; sm_90)
|
||||
docker build --file docker/Dockerfile \
|
||||
--platform linux/arm64 \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg nvcc_threads=4 \
|
||||
--build-arg torch_cuda_arch_list="9.0" \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg buildkite_commit="$BUILDKITE_COMMIT" \
|
||||
--tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 \
|
||||
--target test \
|
||||
--progress plain .
|
||||
|
||||
# push
|
||||
docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64
|
||||
@@ -737,7 +737,7 @@ steps:
|
||||
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
|
||||
env:
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
VARIANT: "rocm723"
|
||||
VARIANT: "rocm722"
|
||||
|
||||
# ROCm Job 6: Build ROCm Release Docker Image
|
||||
- label: ":docker: Build release image - x86_64 - ROCm"
|
||||
|
||||
@@ -35,9 +35,25 @@ export PYTHONPATH=".."
|
||||
# Helper Functions
|
||||
###############################################################################
|
||||
|
||||
report_docker_usage() {
|
||||
echo "--- Docker usage"
|
||||
docker system df || true
|
||||
cleanup_docker() {
|
||||
# Get Docker's root directory
|
||||
docker_root=$(docker info -f '{{.DockerRootDir}}')
|
||||
if [ -z "$docker_root" ]; then
|
||||
echo "Failed to determine Docker root directory."
|
||||
exit 1
|
||||
fi
|
||||
echo "Docker root directory: $docker_root"
|
||||
|
||||
disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
|
||||
threshold=70
|
||||
if [ "$disk_usage" -gt "$threshold" ]; then
|
||||
echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
|
||||
docker image prune -f
|
||||
docker volume prune -f && docker system prune --force --filter "until=72h" --all
|
||||
echo "Docker images and volumes cleanup completed."
|
||||
else
|
||||
echo "Disk usage is below $threshold%. No cleanup needed."
|
||||
fi
|
||||
}
|
||||
|
||||
cleanup_network() {
|
||||
@@ -238,8 +254,8 @@ re_quote_pytest_markers() {
|
||||
echo "--- ROCm info"
|
||||
rocminfo
|
||||
|
||||
# --- Docker status ---
|
||||
report_docker_usage
|
||||
# --- Docker housekeeping ---
|
||||
cleanup_docker
|
||||
|
||||
# --- Pull test image ---
|
||||
echo "--- Pulling container"
|
||||
@@ -248,17 +264,9 @@ container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | hea
|
||||
docker pull "${image_name}"
|
||||
|
||||
remove_docker_container() {
|
||||
# docker run uses --rm, so the container is normally already gone when the
|
||||
# EXIT trap runs. Cleanup is best-effort and must not affect the test result.
|
||||
docker rm -f "${container_name}" >/dev/null 2>&1 || true
|
||||
docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
|
||||
}
|
||||
|
||||
on_exit() {
|
||||
local exit_code=$?
|
||||
remove_docker_container
|
||||
exit "$exit_code"
|
||||
}
|
||||
trap on_exit EXIT
|
||||
trap remove_docker_container EXIT
|
||||
|
||||
# --- Prepare commands ---
|
||||
echo "--- Running container"
|
||||
|
||||
@@ -38,7 +38,7 @@ steps:
|
||||
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 40
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -28,8 +28,7 @@ steps:
|
||||
- pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -46,8 +45,7 @@ steps:
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 80
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
@@ -65,8 +63,7 @@ steps:
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
@@ -85,8 +82,7 @@ steps:
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
@@ -108,8 +104,7 @@ steps:
|
||||
- pytest -v -s tool_use
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
@@ -86,7 +86,7 @@ steps:
|
||||
parallelism: 2
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/
|
||||
- vllm/model_executor/layers/quantization
|
||||
|
||||
@@ -52,7 +52,7 @@ steps:
|
||||
- pytest -v -s v1/test_outputs.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
@@ -58,3 +58,37 @@ steps:
|
||||
device: cpu-small
|
||||
commands:
|
||||
- pytest -v -s models/test_utils.py models/test_vision.py
|
||||
|
||||
- label: Transformers Nightly Models
|
||||
device: h200_35gb
|
||||
key: transformers-nightly-models
|
||||
working_dir: "/vllm-workspace/"
|
||||
optional: true
|
||||
soft_fail: true
|
||||
commands:
|
||||
- pip install --upgrade git+https://github.com/huggingface/transformers
|
||||
- pytest -v -s tests/models/test_initialization.py
|
||||
- pytest -v -s tests/models/test_transformers.py
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
- label: Transformers Backward Compatibility Models Test
|
||||
device: h200_35gb
|
||||
key: transformers-backward-compatibility-models-test
|
||||
working_dir: "/vllm-workspace/"
|
||||
optional: true
|
||||
soft_fail: true
|
||||
commands:
|
||||
- pip install transformers==4.57.5
|
||||
- pytest -v -s tests/models/test_initialization.py
|
||||
- pytest -v -s tests/models/test_transformers.py
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
@@ -50,7 +50,7 @@ steps:
|
||||
mirror:
|
||||
torch_nightly: {}
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
@@ -96,7 +96,7 @@ steps:
|
||||
- pytest -v -s models/language/pooling -m 'not core_model'
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 100
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -15,7 +15,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -33,7 +33,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -50,7 +50,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -118,7 +118,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/test_mapping.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
@@ -32,7 +32,6 @@ steps:
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/transformers_utils/configs/speculators/
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
|
||||
+11
-11
@@ -80,13 +80,13 @@
|
||||
/tests/distributed/test_multi_node_assignment.py @youkaichao
|
||||
/tests/distributed/test_pipeline_parallel.py @youkaichao
|
||||
/tests/distributed/test_same_node.py @youkaichao
|
||||
/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas
|
||||
/tests/evals @mgoin @vadiklyutiy @AndreasKaratzas
|
||||
/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas
|
||||
/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche
|
||||
/tests/evals @mgoin @vadiklyutiy
|
||||
/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye
|
||||
/tests/kernels/ir @ProExpertProg @tjtanaa
|
||||
/tests/models @DarkLight1337 @ywang96 @AndreasKaratzas
|
||||
/tests/models @DarkLight1337 @ywang96
|
||||
/tests/multimodal @DarkLight1337 @ywang96 @NickLucche
|
||||
/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas
|
||||
/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye
|
||||
/tests/test_inputs.py @DarkLight1337 @ywang96
|
||||
/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
|
||||
/tests/v1/structured_output @mgoin @russellb @aarnphm
|
||||
@@ -171,20 +171,20 @@ mkdocs.yaml @hmellor
|
||||
|
||||
# ROCm related: specify owner with write access to notify AMD folks for careful code review
|
||||
/vllm/**/*rocm* @tjtanaa @dllehr-amd
|
||||
/docker/Dockerfile.rocm* @tjtanaa @dllehr-amd @AndreasKaratzas
|
||||
/docker/Dockerfile.rocm* @tjtanaa @dllehr-amd
|
||||
/vllm/v1/attention/backends/rocm*.py @tjtanaa @dllehr-amd
|
||||
/vllm/v1/attention/backends/mla/rocm*.py @tjtanaa @dllehr-amd
|
||||
/vllm/v1/attention/ops/rocm*.py @tjtanaa @dllehr-amd
|
||||
/vllm/model_executor/layers/fused_moe/rocm*.py @tjtanaa @dllehr-amd
|
||||
/csrc/rocm @tjtanaa @dllehr-amd
|
||||
/requirements/*rocm* @tjtanaa @AndreasKaratzas
|
||||
/tests/**/*rocm* @tjtanaa @AndreasKaratzas
|
||||
/requirements/*rocm* @tjtanaa
|
||||
/tests/**/*rocm* @tjtanaa
|
||||
/docs/**/*rocm* @tjtanaa
|
||||
/vllm/**/*quark* @tjtanaa
|
||||
/tests/**/*quark* @tjtanaa @AndreasKaratzas
|
||||
/tests/**/*quark* @tjtanaa
|
||||
/docs/**/*quark* @tjtanaa
|
||||
/vllm/**/*aiter* @tjtanaa @AndreasKaratzas
|
||||
/tests/**/*aiter* @tjtanaa @AndreasKaratzas
|
||||
/vllm/**/*aiter* @tjtanaa
|
||||
/tests/**/*aiter* @tjtanaa
|
||||
|
||||
# TPU
|
||||
/vllm/v1/worker/tpu* @NickLucche
|
||||
|
||||
@@ -103,19 +103,6 @@ pull_request_rules:
|
||||
add:
|
||||
- frontend
|
||||
|
||||
- name: label-rust
|
||||
description: Automatically apply rust label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=(?i)rust
|
||||
- title~=(?i)rust
|
||||
- title~=(?i)vllm-rs
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- rust
|
||||
|
||||
- name: label-llama
|
||||
description: Automatically apply llama label
|
||||
conditions:
|
||||
|
||||
+9
-27
@@ -305,10 +305,14 @@ endif()
|
||||
#
|
||||
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/mamba/mamba_ssm/selective_scan_fwd.cu"
|
||||
"csrc/cache_kernels.cu"
|
||||
"csrc/cache_kernels_fused.cu"
|
||||
"csrc/attention/paged_attention_v1.cu"
|
||||
"csrc/attention/paged_attention_v2.cu"
|
||||
"csrc/attention/merge_attn_states.cu"
|
||||
"csrc/sampler.cu"
|
||||
"csrc/topk.cu"
|
||||
"csrc/cuda_view.cu"
|
||||
"csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
|
||||
"csrc/quantization/activation_kernels.cu"
|
||||
@@ -365,30 +369,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
# are not supported by Machete yet.
|
||||
|
||||
# marlin arches for fp16 output
|
||||
# Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0;
|
||||
# fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8).
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
|
||||
# marlin has limited support for turing
|
||||
cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
|
||||
# marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
|
||||
# marlin arches for fp8 input
|
||||
# - sm80 doesn't support fp8 computation
|
||||
# - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
|
||||
# so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0;12.1" "${CUDA_ARCHS}")
|
||||
# marlin arches for other files
|
||||
cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
|
||||
|
||||
@@ -643,11 +633,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
"csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
|
||||
"csrc/libtorch_stable/layernorm_kernels.cu"
|
||||
"csrc/libtorch_stable/layernorm_quant_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
|
||||
"csrc/libtorch_stable/attention/merge_attn_states.cu"
|
||||
"csrc/libtorch_stable/sampler.cu"
|
||||
"csrc/libtorch_stable/topk.cu"
|
||||
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu")
|
||||
"csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu")
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC
|
||||
@@ -1135,11 +1121,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
# moe marlin arches
|
||||
# note that we always set `use_atomic_add=False` for moe marlin now,
|
||||
# so we don't need 9.0 for bf16 atomicAdd PTX
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
|
||||
# moe marlin has limited support for turing
|
||||
cuda_archs_loose_intersection(MARLIN_MOE_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
|
||||
# moe marlin arches for fp8 input
|
||||
|
||||
@@ -10,7 +10,6 @@ from transformers import AutoConfig
|
||||
|
||||
from vllm.model_executor.layers.fused_moe import fused_topk
|
||||
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
|
||||
MoEPermuteScratch,
|
||||
moe_permute,
|
||||
moe_unpermute,
|
||||
)
|
||||
@@ -55,15 +54,6 @@ def benchmark_permute(
|
||||
topk_weights, topk_ids, token_expert_indices = fused_topk(
|
||||
qhidden_states, input_gating, topk, False
|
||||
)
|
||||
scratch = MoEPermuteScratch(
|
||||
max_num_tokens=num_tokens,
|
||||
topk=topk,
|
||||
num_experts=num_experts,
|
||||
num_local_experts=num_experts,
|
||||
device=qhidden_states.device,
|
||||
hidden_size=hidden_size,
|
||||
hidden_dtype=qhidden_states.dtype,
|
||||
)
|
||||
|
||||
def prepare(i: int):
|
||||
input_gating.copy_(gating_output[i])
|
||||
@@ -75,7 +65,6 @@ def benchmark_permute(
|
||||
topk_ids=topk_ids,
|
||||
n_expert=num_experts,
|
||||
expert_map=None,
|
||||
scratch=scratch,
|
||||
)
|
||||
|
||||
# JIT compilation & warmup
|
||||
@@ -134,15 +123,6 @@ def benchmark_unpermute(
|
||||
topk_weights, topk_ids, token_expert_indices = fused_topk(
|
||||
qhidden_states, input_gating, topk, False
|
||||
)
|
||||
scratch = MoEPermuteScratch(
|
||||
max_num_tokens=num_tokens,
|
||||
topk=topk,
|
||||
num_experts=num_experts,
|
||||
num_local_experts=num_experts,
|
||||
device=qhidden_states.device,
|
||||
hidden_size=hidden_size,
|
||||
hidden_dtype=qhidden_states.dtype,
|
||||
)
|
||||
|
||||
def prepare():
|
||||
(
|
||||
@@ -157,7 +137,6 @@ def benchmark_unpermute(
|
||||
topk_ids=topk_ids,
|
||||
n_expert=num_experts,
|
||||
expert_map=None,
|
||||
scratch=scratch,
|
||||
)
|
||||
# convert to fp16/bf16 as gemm output
|
||||
return (
|
||||
|
||||
+49
-56
@@ -1,14 +1,14 @@
|
||||
#include <optional>
|
||||
#include <torch/all.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <algorithm>
|
||||
#include <limits>
|
||||
|
||||
#include "../torch_utils.h"
|
||||
#include "attention_dtypes.h"
|
||||
#include "attention_utils.cuh"
|
||||
#include "../quantization/w8a8/fp8/common.cuh"
|
||||
#include "../dispatch_utils.h"
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
|
||||
#include "../../attention/attention_dtypes.h"
|
||||
#include "../../attention/attention_utils.cuh"
|
||||
#include "../../quantization/w8a8/fp8/common.cuh"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -196,17 +196,17 @@ __global__ void merge_attn_states_kernel(
|
||||
// The following macro is used to dispatch the conversion function based on
|
||||
// the output data type. The FN is a macro that calls a function with
|
||||
// template<typename scalar_t>.
|
||||
#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn) \
|
||||
{ \
|
||||
if (scalar_dtype == torch::headeronly::ScalarType::Float) { \
|
||||
fn(float); \
|
||||
} else if (scalar_dtype == torch::headeronly::ScalarType::Half) { \
|
||||
fn(uint16_t); \
|
||||
} else if (scalar_dtype == torch::headeronly::ScalarType::BFloat16) { \
|
||||
fn(__nv_bfloat16); \
|
||||
} else { \
|
||||
STD_TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
|
||||
} \
|
||||
#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn) \
|
||||
{ \
|
||||
if (scalar_dtype == at::ScalarType::Float) { \
|
||||
fn(float); \
|
||||
} else if (scalar_dtype == at::ScalarType::Half) { \
|
||||
fn(uint16_t); \
|
||||
} else if (scalar_dtype == at::ScalarType::BFloat16) { \
|
||||
fn(__nv_bfloat16); \
|
||||
} else { \
|
||||
TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
|
||||
} \
|
||||
}
|
||||
|
||||
#define LAUNCH_MERGE_ATTN_STATES(scalar_t, output_t, NUM_THREADS, \
|
||||
@@ -245,14 +245,11 @@ __global__ void merge_attn_states_kernel(
|
||||
*/
|
||||
template <typename scalar_t>
|
||||
void merge_attn_states_launcher(
|
||||
torch::stable::Tensor& output,
|
||||
std::optional<torch::stable::Tensor> output_lse,
|
||||
const torch::stable::Tensor& prefix_output,
|
||||
const torch::stable::Tensor& prefix_lse,
|
||||
const torch::stable::Tensor& suffix_output,
|
||||
const torch::stable::Tensor& suffix_lse,
|
||||
torch::Tensor& output, std::optional<torch::Tensor> output_lse,
|
||||
const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
|
||||
const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
|
||||
const std::optional<int64_t> prefill_tokens_with_context,
|
||||
const std::optional<torch::stable::Tensor>& output_scale) {
|
||||
const std::optional<torch::Tensor>& output_scale) {
|
||||
constexpr uint NUM_THREADS = 128;
|
||||
const uint num_tokens = output.size(0);
|
||||
const uint num_heads = output.size(1);
|
||||
@@ -261,23 +258,23 @@ void merge_attn_states_launcher(
|
||||
const uint output_head_stride = output.stride(1);
|
||||
// Thread mapping is based on input BF16 pack_size
|
||||
const uint pack_size = 16 / sizeof(scalar_t);
|
||||
STD_TORCH_CHECK(head_size % pack_size == 0,
|
||||
"headsize must be multiple of pack_size:", pack_size);
|
||||
TORCH_CHECK(head_size % pack_size == 0,
|
||||
"headsize must be multiple of pack_size:", pack_size);
|
||||
|
||||
const uint prefix_num_tokens =
|
||||
prefill_tokens_with_context.has_value()
|
||||
? static_cast<uint>(prefill_tokens_with_context.value())
|
||||
: num_tokens;
|
||||
STD_TORCH_CHECK(prefix_num_tokens <= num_tokens,
|
||||
"prefix_num_tokens must be <= num_tokens");
|
||||
TORCH_CHECK(prefix_num_tokens <= num_tokens,
|
||||
"prefix_num_tokens must be <= num_tokens");
|
||||
|
||||
float* output_lse_ptr = nullptr;
|
||||
if (output_lse.has_value()) {
|
||||
output_lse_ptr = output_lse.value().mutable_data_ptr<float>();
|
||||
output_lse_ptr = output_lse.value().data_ptr<float>();
|
||||
}
|
||||
float* output_scale_ptr = nullptr;
|
||||
if (output_scale.has_value()) {
|
||||
output_scale_ptr = output_scale.value().mutable_data_ptr<float>();
|
||||
output_scale_ptr = output_scale.value().data_ptr<float>();
|
||||
}
|
||||
// Process one pack elements per thread. for float, the
|
||||
// pack_size is 4 for half/bf16, the pack_size is 8.
|
||||
@@ -287,15 +284,14 @@ void merge_attn_states_launcher(
|
||||
dim3 block(NUM_THREADS);
|
||||
dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS);
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
prefix_output.get_device_index());
|
||||
auto stream = get_current_cuda_stream();
|
||||
const c10::cuda::OptionalCUDAGuard device_guard(prefix_output.device());
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
if (output_scale.has_value()) {
|
||||
// FP8 output path - dispatch on output FP8 type
|
||||
VLLM_STABLE_DISPATCH_FP8_TYPES(
|
||||
output.scalar_type(), "merge_attn_states_fp8",
|
||||
[&] { LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true); });
|
||||
VLLM_DISPATCH_FP8_TYPES(output.scalar_type(), "merge_attn_states_fp8", [&] {
|
||||
LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true);
|
||||
});
|
||||
} else {
|
||||
// Original BF16/FP16/FP32 output path
|
||||
LAUNCH_MERGE_ATTN_STATES(scalar_t, scalar_t, NUM_THREADS, false);
|
||||
@@ -309,29 +305,26 @@ void merge_attn_states_launcher(
|
||||
suffix_lse, prefill_tokens_with_context, output_scale); \
|
||||
}
|
||||
|
||||
void merge_attn_states(
|
||||
torch::stable::Tensor& output,
|
||||
std::optional<torch::stable::Tensor> output_lse,
|
||||
const torch::stable::Tensor& prefix_output,
|
||||
const torch::stable::Tensor& prefix_lse,
|
||||
const torch::stable::Tensor& suffix_output,
|
||||
const torch::stable::Tensor& suffix_lse,
|
||||
const std::optional<int64_t> prefill_tokens_with_context,
|
||||
const std::optional<torch::stable::Tensor>& output_scale) {
|
||||
void merge_attn_states(torch::Tensor& output,
|
||||
std::optional<torch::Tensor> output_lse,
|
||||
const torch::Tensor& prefix_output,
|
||||
const torch::Tensor& prefix_lse,
|
||||
const torch::Tensor& suffix_output,
|
||||
const torch::Tensor& suffix_lse,
|
||||
std::optional<int64_t> prefill_tokens_with_context,
|
||||
const std::optional<torch::Tensor>& output_scale) {
|
||||
if (output_scale.has_value()) {
|
||||
STD_TORCH_CHECK(
|
||||
output.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
|
||||
output.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fnuz,
|
||||
"output must be FP8 when output_scale is provided, got: ",
|
||||
output.scalar_type());
|
||||
TORCH_CHECK(output.scalar_type() == at::ScalarType::Float8_e4m3fn ||
|
||||
output.scalar_type() == at::ScalarType::Float8_e4m3fnuz,
|
||||
"output must be FP8 when output_scale is provided, got: ",
|
||||
output.scalar_type());
|
||||
} else {
|
||||
STD_TORCH_CHECK(
|
||||
output.scalar_type() == prefix_output.scalar_type(), "output dtype (",
|
||||
output.scalar_type(), ") must match prefix_output dtype (",
|
||||
prefix_output.scalar_type(), ") when output_scale is not set");
|
||||
TORCH_CHECK(output.scalar_type() == prefix_output.scalar_type(),
|
||||
"output dtype (", output.scalar_type(),
|
||||
") must match prefix_output dtype (",
|
||||
prefix_output.scalar_type(), ") when output_scale is not set");
|
||||
}
|
||||
// Always dispatch on prefix_output (input) dtype
|
||||
DISPATCH_BY_SCALAR_DTYPE(prefix_output.scalar_type(),
|
||||
DISPATCH_BY_SCALAR_DTYPE(prefix_output.dtype(),
|
||||
CALL_MERGE_ATTN_STATES_LAUNCHER);
|
||||
}
|
||||
@@ -164,17 +164,6 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
|
||||
|
||||
#endif
|
||||
|
||||
// Attention kernels (shared CUDA/ROCm)
|
||||
void merge_attn_states(
|
||||
torch::stable::Tensor& output,
|
||||
std::optional<torch::stable::Tensor> output_lse,
|
||||
const torch::stable::Tensor& prefix_output,
|
||||
const torch::stable::Tensor& prefix_lse,
|
||||
const torch::stable::Tensor& suffix_output,
|
||||
const torch::stable::Tensor& suffix_lse,
|
||||
const std::optional<int64_t> prefill_tokens_with_context,
|
||||
const std::optional<torch::stable::Tensor>& output_scale = std::nullopt);
|
||||
|
||||
torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
|
||||
bool inplace);
|
||||
|
||||
@@ -231,48 +220,6 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
|
||||
torch::stable::Tensor& position_ids,
|
||||
int64_t forced_token_heads_per_warp);
|
||||
|
||||
// Sampler kernels (shared CUDA/ROCm)
|
||||
void apply_repetition_penalties_(
|
||||
torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
|
||||
const torch::stable::Tensor& output_mask,
|
||||
const torch::stable::Tensor& repetition_penalties);
|
||||
|
||||
void top_k_per_row_prefill(const torch::stable::Tensor& logits,
|
||||
const torch::stable::Tensor& rowStarts,
|
||||
const torch::stable::Tensor& rowEnds,
|
||||
torch::stable::Tensor& indices, int64_t numRows,
|
||||
int64_t stride0, int64_t stride1, int64_t topK);
|
||||
|
||||
void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
|
||||
const torch::stable::Tensor& seqLens,
|
||||
torch::stable::Tensor& indices, int64_t numRows,
|
||||
int64_t stride0, int64_t stride1, int64_t topK);
|
||||
|
||||
void persistent_topk(const torch::stable::Tensor& logits,
|
||||
const torch::stable::Tensor& lengths,
|
||||
torch::stable::Tensor& output,
|
||||
torch::stable::Tensor& workspace, int64_t k,
|
||||
int64_t max_seq_len);
|
||||
|
||||
void selective_scan_fwd(
|
||||
const torch::stable::Tensor& u, const torch::stable::Tensor& delta,
|
||||
const torch::stable::Tensor& A, const torch::stable::Tensor& B,
|
||||
const torch::stable::Tensor& C,
|
||||
const std::optional<torch::stable::Tensor>& D_,
|
||||
const std::optional<torch::stable::Tensor>& z_,
|
||||
const std::optional<torch::stable::Tensor>& delta_bias_,
|
||||
bool delta_softplus,
|
||||
const std::optional<torch::stable::Tensor>& query_start_loc,
|
||||
const std::optional<torch::stable::Tensor>& cache_indices,
|
||||
const std::optional<torch::stable::Tensor>& has_initial_state,
|
||||
const torch::stable::Tensor& ssm_states, int64_t null_block_id,
|
||||
int64_t block_size,
|
||||
const std::optional<torch::stable::Tensor>& block_idx_first_scheduled_token,
|
||||
const std::optional<torch::stable::Tensor>& block_idx_last_scheduled_token,
|
||||
const std::optional<torch::stable::Tensor>& initial_state_idx,
|
||||
const std::optional<torch::stable::Tensor>& cu_chunk_seqlen,
|
||||
const std::optional<torch::stable::Tensor>& last_chunk_indices);
|
||||
|
||||
// Activation kernels (shared CUDA/ROCm)
|
||||
void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
|
||||
void silu_and_mul_clamp(torch::stable::Tensor& out,
|
||||
|
||||
@@ -263,20 +263,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
|
||||
#endif
|
||||
|
||||
// Merge attn states
|
||||
// Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
|
||||
// can be used to combine partial attention results (in the split-KV case)
|
||||
ops.def(
|
||||
"merge_attn_states("
|
||||
" Tensor! output,"
|
||||
" Tensor!? output_lse,"
|
||||
" Tensor prefix_output,"
|
||||
" Tensor prefix_lse,"
|
||||
" Tensor suffix_output,"
|
||||
" Tensor suffix_lse,"
|
||||
" int!? prefill_tokens_with_context,"
|
||||
" Tensor? output_scale=None) -> ()");
|
||||
|
||||
// Hadamard transforms
|
||||
// conditionally compiled so impl registration is in source file
|
||||
ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
|
||||
@@ -333,26 +319,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"bool is_neox, Tensor position_ids, "
|
||||
"int forced_token_heads_per_warp=-1) -> ()");
|
||||
|
||||
// Apply repetition penalties to logits in-place.
|
||||
ops.def(
|
||||
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
|
||||
"Tensor output_mask, Tensor repetition_penalties) -> ()");
|
||||
|
||||
// Optimized top-k per row operations.
|
||||
ops.def(
|
||||
"top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
|
||||
"Tensor! indices, int numRows, int stride0, "
|
||||
"int stride1, int topK) -> ()");
|
||||
|
||||
ops.def(
|
||||
"top_k_per_row_decode(Tensor logits, int next_n, "
|
||||
"Tensor seq_lens, Tensor! indices, "
|
||||
"int numRows, int stride0, int stride1, int topK) -> ()");
|
||||
|
||||
ops.def(
|
||||
"persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
|
||||
"Tensor workspace, int k, int max_seq_len) -> ()");
|
||||
|
||||
// Activation ops
|
||||
// Activation function used in SwiGLU.
|
||||
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
|
||||
@@ -456,24 +422,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"int type, SymInt row, SymInt tokens) -> Tensor");
|
||||
|
||||
ops.def("ggml_moe_get_block_size(int type) -> int");
|
||||
|
||||
// Mamba selective scan kernel
|
||||
ops.def(
|
||||
"selective_scan_fwd(Tensor! u, Tensor! delta,"
|
||||
"Tensor! A, Tensor! B, Tensor! C,"
|
||||
"Tensor? D_, Tensor!? z_, Tensor? delta_bias_,"
|
||||
"bool delta_softplus,"
|
||||
"Tensor? query_start_loc,"
|
||||
"Tensor? cache_indices,"
|
||||
"Tensor? has_initial_state,"
|
||||
"Tensor! ssm_states,"
|
||||
"int null_block_id,"
|
||||
"int block_size,"
|
||||
"Tensor? block_idx_first_scheduled_token,"
|
||||
"Tensor? block_idx_last_scheduled_token,"
|
||||
"Tensor? initial_state_idx,"
|
||||
"Tensor? cu_chunk_seqlen,"
|
||||
"Tensor? last_chunk_indices) -> ()");
|
||||
}
|
||||
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
@@ -521,8 +469,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
// files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
|
||||
#endif
|
||||
|
||||
ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states));
|
||||
|
||||
// Layernorm kernels (shared CUDA/ROCm)
|
||||
ops.impl("rms_norm", TORCH_BOX(&rms_norm));
|
||||
ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm));
|
||||
@@ -541,13 +487,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
|
||||
ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
|
||||
|
||||
// Sampler kernels (shared CUDA/ROCm)
|
||||
ops.impl("apply_repetition_penalties_",
|
||||
TORCH_BOX(&apply_repetition_penalties_));
|
||||
ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill));
|
||||
ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode));
|
||||
ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
|
||||
|
||||
// Activation kernels (shared CUDA/ROCm)
|
||||
ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
|
||||
ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
|
||||
@@ -580,7 +519,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("ggml_mul_mat_a8", TORCH_BOX(&ggml_mul_mat_a8));
|
||||
ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8));
|
||||
ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec));
|
||||
ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd));
|
||||
}
|
||||
|
||||
// These capability-check functions take only primitive args (no tensors), so
|
||||
|
||||
@@ -12,9 +12,6 @@
|
||||
#include <hip/hip_bf16.h>
|
||||
#endif
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
#include <torch/headeronly/util/Half.h>
|
||||
#include <torch/headeronly/util/BFloat16.h>
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct SSMParamsBase {
|
||||
@@ -162,8 +159,8 @@ struct Converter{
|
||||
};
|
||||
|
||||
template<int N>
|
||||
struct Converter<torch::headeronly::Half, N>{
|
||||
static inline __device__ void to_float(const torch::headeronly::Half (&src)[N], float (&dst)[N]) {
|
||||
struct Converter<at::Half, N>{
|
||||
static inline __device__ void to_float(const at::Half (&src)[N], float (&dst)[N]) {
|
||||
static_assert(N % 2 == 0);
|
||||
auto &src2 = reinterpret_cast<const half2 (&)[N / 2]>(src);
|
||||
auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
|
||||
@@ -174,8 +171,8 @@ struct Converter<torch::headeronly::Half, N>{
|
||||
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
template<int N>
|
||||
struct Converter<torch::headeronly::BFloat16, N>{
|
||||
static inline __device__ void to_float(const torch::headeronly::BFloat16 (&src)[N], float (&dst)[N]) {
|
||||
struct Converter<at::BFloat16, N>{
|
||||
static inline __device__ void to_float(const at::BFloat16 (&src)[N], float (&dst)[N]) {
|
||||
static_assert(N % 2 == 0);
|
||||
auto &src2 = reinterpret_cast<const nv_bfloat162 (&)[N / 2]>(src);
|
||||
auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
|
||||
+111
-102
@@ -1,9 +1,18 @@
|
||||
// clang-format off
|
||||
// adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/selective_scan_fwd_kernel.cuh
|
||||
#include "../torch_utils.h"
|
||||
#include <torch/csrc/stable/macros.h>
|
||||
#include <torch/all.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include "selective_scan.h"
|
||||
|
||||
#include <c10/util/BFloat16.h>
|
||||
#include <c10/util/Half.h>
|
||||
#ifdef USE_ROCM
|
||||
#include <c10/hip/HIPException.h> // For C10_HIP_CHECK and C10_HIP_KERNEL_LAUNCH_CHECK
|
||||
#else
|
||||
#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#include <cub/block/block_load.cuh>
|
||||
#include <cub/block/block_store.cuh>
|
||||
@@ -407,15 +416,15 @@ void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) {
|
||||
auto kernel = &selective_scan_fwd_kernel<Ktraits>;
|
||||
if (kSmemSize >= 48 * 1024) {
|
||||
#ifdef USE_ROCM
|
||||
STD_CUDA_CHECK(hipFuncSetAttribute(
|
||||
C10_HIP_CHECK(hipFuncSetAttribute(
|
||||
reinterpret_cast<const void*>(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
|
||||
#else
|
||||
STD_CUDA_CHECK(cudaFuncSetAttribute(
|
||||
C10_CUDA_CHECK(cudaFuncSetAttribute(
|
||||
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
|
||||
#endif
|
||||
}
|
||||
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
|
||||
STD_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -453,46 +462,46 @@ void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream) {
|
||||
#endif
|
||||
}
|
||||
|
||||
template void selective_scan_fwd_cuda<torch::headeronly::BFloat16, float, torch::headeronly::BFloat16>(SSMParamsBase ¶ms, cudaStream_t stream);
|
||||
template void selective_scan_fwd_cuda<torch::headeronly::BFloat16, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
|
||||
template void selective_scan_fwd_cuda<torch::headeronly::Half, float, torch::headeronly::Half>(SSMParamsBase ¶ms, cudaStream_t stream);
|
||||
template void selective_scan_fwd_cuda<torch::headeronly::Half, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
|
||||
template void selective_scan_fwd_cuda<at::BFloat16, float, at::BFloat16>(SSMParamsBase ¶ms, cudaStream_t stream);
|
||||
template void selective_scan_fwd_cuda<at::BFloat16, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
|
||||
template void selective_scan_fwd_cuda<at::Half, float, at::Half>(SSMParamsBase ¶ms, cudaStream_t stream);
|
||||
template void selective_scan_fwd_cuda<at::Half, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
|
||||
template void selective_scan_fwd_cuda<float, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
|
||||
|
||||
#define CHECK_SHAPE(x, ...) STD_TORCH_CHECK(x.sizes().equals(torch::headeronly::IntHeaderOnlyArrayRef({__VA_ARGS__})), #x " must have shape (" #__VA_ARGS__ ")")
|
||||
#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
|
||||
|
||||
#define DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, STYPE, NAME, ...) \
|
||||
if (ITYPE == torch::headeronly::ScalarType::Half) { \
|
||||
using input_t = torch::headeronly::Half; \
|
||||
if (ITYPE == at::ScalarType::Half) { \
|
||||
using input_t = at::Half; \
|
||||
using weight_t = float; \
|
||||
if (STYPE == torch::headeronly::ScalarType::Half) { \
|
||||
using state_t = torch::headeronly::Half; \
|
||||
if (STYPE == at::ScalarType::Half) { \
|
||||
using state_t = at::Half; \
|
||||
__VA_ARGS__(); \
|
||||
} else if (STYPE == torch::headeronly::ScalarType::Float) { \
|
||||
} else if (STYPE == at::ScalarType::Float) { \
|
||||
using state_t = float; \
|
||||
__VA_ARGS__(); \
|
||||
} else { \
|
||||
STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \
|
||||
AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \
|
||||
} \
|
||||
} else if (ITYPE == torch::headeronly::ScalarType::BFloat16) { \
|
||||
using input_t = torch::headeronly::BFloat16; \
|
||||
} else if (ITYPE == at::ScalarType::BFloat16) { \
|
||||
using input_t = at::BFloat16; \
|
||||
using weight_t = float; \
|
||||
if (STYPE == torch::headeronly::ScalarType::BFloat16) { \
|
||||
using state_t = torch::headeronly::BFloat16; \
|
||||
if (STYPE == at::ScalarType::BFloat16) { \
|
||||
using state_t = at::BFloat16; \
|
||||
__VA_ARGS__(); \
|
||||
} else if (STYPE == torch::headeronly::ScalarType::Float) { \
|
||||
} else if (STYPE == at::ScalarType::Float) { \
|
||||
using state_t = float; \
|
||||
__VA_ARGS__(); \
|
||||
} else { \
|
||||
STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \
|
||||
AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \
|
||||
} \
|
||||
} else if (ITYPE == torch::headeronly::ScalarType::Float) { \
|
||||
} else if (ITYPE == at::ScalarType::Float) { \
|
||||
using input_t = float; \
|
||||
using weight_t = float; \
|
||||
using state_t = float; \
|
||||
__VA_ARGS__(); \
|
||||
} else { \
|
||||
STD_TORCH_CHECK(false, #NAME " not implemented for input type '", ITYPE, "'"); \
|
||||
AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
|
||||
}
|
||||
|
||||
|
||||
@@ -509,30 +518,30 @@ void set_ssm_params_fwd(SSMParamsBase ¶ms,
|
||||
const bool is_variable_B,
|
||||
const bool is_variable_C,
|
||||
// device pointers
|
||||
const torch::stable::Tensor u,
|
||||
const torch::stable::Tensor delta,
|
||||
const torch::stable::Tensor A,
|
||||
const torch::stable::Tensor B,
|
||||
const torch::stable::Tensor C,
|
||||
const torch::stable::Tensor out,
|
||||
const torch::stable::Tensor z,
|
||||
const torch::stable::Tensor out_z,
|
||||
const std::optional<torch::stable::Tensor>& D,
|
||||
const std::optional<torch::stable::Tensor>& delta_bias,
|
||||
const torch::stable::Tensor ssm_states,
|
||||
const torch::Tensor u,
|
||||
const torch::Tensor delta,
|
||||
const torch::Tensor A,
|
||||
const torch::Tensor B,
|
||||
const torch::Tensor C,
|
||||
const torch::Tensor out,
|
||||
const torch::Tensor z,
|
||||
const torch::Tensor out_z,
|
||||
const std::optional<at::Tensor>& D,
|
||||
const std::optional<at::Tensor>& delta_bias,
|
||||
const torch::Tensor ssm_states,
|
||||
bool has_z,
|
||||
bool delta_softplus,
|
||||
const std::optional<torch::stable::Tensor>& query_start_loc,
|
||||
const std::optional<torch::stable::Tensor>& cache_indices,
|
||||
const std::optional<torch::stable::Tensor>& has_initial_state,
|
||||
const std::optional<at::Tensor>& query_start_loc,
|
||||
const std::optional<at::Tensor>& cache_indices,
|
||||
const std::optional<at::Tensor>& has_initial_state,
|
||||
bool varlen,
|
||||
int64_t null_block_id,
|
||||
int64_t block_size,
|
||||
const std::optional<torch::stable::Tensor> &block_idx_first_scheduled_token,
|
||||
const std::optional<torch::stable::Tensor> &block_idx_last_scheduled_token,
|
||||
const std::optional<torch::stable::Tensor> &initial_state_idx,
|
||||
const std::optional<torch::stable::Tensor> &cu_chunk_seqlen,
|
||||
const std::optional<torch::stable::Tensor> &last_chunk_indices) {
|
||||
const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
|
||||
const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
|
||||
const std::optional<torch::Tensor> &initial_state_idx,
|
||||
const std::optional<torch::Tensor> &cu_chunk_seqlen,
|
||||
const std::optional<torch::Tensor> &last_chunk_indices) {
|
||||
|
||||
// Reset the parameters
|
||||
memset(¶ms, 0, sizeof(params));
|
||||
@@ -645,45 +654,45 @@ void set_ssm_params_fwd(SSMParamsBase ¶ms,
|
||||
}
|
||||
}
|
||||
|
||||
void selective_scan_fwd(const torch::stable::Tensor &u, const torch::stable::Tensor &delta,
|
||||
const torch::stable::Tensor &A, const torch::stable::Tensor &B, const torch::stable::Tensor &C,
|
||||
const std::optional<torch::stable::Tensor> &D_,
|
||||
const std::optional<torch::stable::Tensor> &z_,
|
||||
const std::optional<torch::stable::Tensor> &delta_bias_,
|
||||
void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
|
||||
const torch::Tensor &A, const torch::Tensor &B, const torch::Tensor &C,
|
||||
const std::optional<torch::Tensor> &D_,
|
||||
const std::optional<torch::Tensor> &z_,
|
||||
const std::optional<torch::Tensor> &delta_bias_,
|
||||
bool delta_softplus,
|
||||
const std::optional<torch::stable::Tensor> &query_start_loc,
|
||||
const std::optional<torch::stable::Tensor> &cache_indices,
|
||||
const std::optional<torch::stable::Tensor> &has_initial_state,
|
||||
const torch::stable::Tensor &ssm_states,
|
||||
const std::optional<torch::Tensor> &query_start_loc,
|
||||
const std::optional<torch::Tensor> &cache_indices,
|
||||
const std::optional<torch::Tensor> &has_initial_state,
|
||||
const torch::Tensor &ssm_states,
|
||||
// used to identify padding entries if cache_indices provided
|
||||
// in case of padding, the kernel will return early
|
||||
int64_t null_block_id,
|
||||
int64_t block_size,
|
||||
const std::optional<torch::stable::Tensor> &block_idx_first_scheduled_token,
|
||||
const std::optional<torch::stable::Tensor> &block_idx_last_scheduled_token,
|
||||
const std::optional<torch::stable::Tensor> &initial_state_idx,
|
||||
const std::optional<torch::stable::Tensor> &cu_chunk_seqlen,
|
||||
const std::optional<torch::stable::Tensor> &last_chunk_indices) {
|
||||
const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
|
||||
const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
|
||||
const std::optional<torch::Tensor> &initial_state_idx,
|
||||
const std::optional<torch::Tensor> &cu_chunk_seqlen,
|
||||
const std::optional<torch::Tensor> &last_chunk_indices) {
|
||||
auto input_type = u.scalar_type();
|
||||
auto weight_type = A.scalar_type();
|
||||
STD_TORCH_CHECK(input_type == torch::headeronly::ScalarType::Float || input_type == torch::headeronly::ScalarType::Half || input_type == torch::headeronly::ScalarType::BFloat16);
|
||||
STD_TORCH_CHECK(weight_type == torch::headeronly::ScalarType::Float);
|
||||
TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
|
||||
TORCH_CHECK(weight_type == at::ScalarType::Float);
|
||||
|
||||
const bool is_variable_B = B.dim() >= 3;
|
||||
const bool is_variable_C = C.dim() >= 3;
|
||||
|
||||
STD_TORCH_CHECK(delta.scalar_type() == input_type);
|
||||
STD_TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type));
|
||||
STD_TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type));
|
||||
TORCH_CHECK(delta.scalar_type() == input_type);
|
||||
TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type));
|
||||
TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type));
|
||||
|
||||
STD_TORCH_CHECK(u.is_cuda());
|
||||
STD_TORCH_CHECK(delta.is_cuda());
|
||||
STD_TORCH_CHECK(A.is_cuda());
|
||||
STD_TORCH_CHECK(B.is_cuda());
|
||||
STD_TORCH_CHECK(C.is_cuda());
|
||||
TORCH_CHECK(u.is_cuda());
|
||||
TORCH_CHECK(delta.is_cuda());
|
||||
TORCH_CHECK(A.is_cuda());
|
||||
TORCH_CHECK(B.is_cuda());
|
||||
TORCH_CHECK(C.is_cuda());
|
||||
|
||||
STD_TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
|
||||
STD_TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
|
||||
TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
|
||||
TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
|
||||
|
||||
const auto sizes = u.sizes();
|
||||
const bool varlen = query_start_loc.has_value();
|
||||
@@ -693,7 +702,7 @@ void selective_scan_fwd(const torch::stable::Tensor &u, const torch::stable::Ten
|
||||
const int dstate = A.size(1);
|
||||
const int n_groups = varlen ? B.size(0) : B.size(1);
|
||||
|
||||
STD_TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256");
|
||||
TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256");
|
||||
|
||||
if (varlen) {
|
||||
CHECK_SHAPE(u, dim, seqlen);
|
||||
@@ -703,94 +712,94 @@ void selective_scan_fwd(const torch::stable::Tensor &u, const torch::stable::Ten
|
||||
CHECK_SHAPE(delta, batch_size, dim, seqlen);
|
||||
}
|
||||
CHECK_SHAPE(A, dim, dstate);
|
||||
STD_TORCH_CHECK(is_variable_B, "is_variable_B = False is disabled in favor of reduced binary size");
|
||||
TORCH_CHECK(is_variable_B, "is_variable_B = False is disabled in favor of reduced binary size")
|
||||
if (varlen) {
|
||||
CHECK_SHAPE(B, n_groups, dstate, seqlen);
|
||||
} else {
|
||||
CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
|
||||
CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
|
||||
}
|
||||
STD_TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
|
||||
TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
|
||||
|
||||
STD_TORCH_CHECK(is_variable_C, "is_variable_C = False is disabled in favor of reduced binary size");
|
||||
TORCH_CHECK(is_variable_C, "is_variable_C = False is disabled in favor of reduced binary size")
|
||||
if (varlen) {
|
||||
CHECK_SHAPE(C, n_groups, dstate, seqlen);
|
||||
} else {
|
||||
CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
|
||||
CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
|
||||
}
|
||||
STD_TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
|
||||
TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
|
||||
|
||||
if (D_.has_value()) {
|
||||
auto D = D_.value();
|
||||
STD_TORCH_CHECK(D.scalar_type() == torch::headeronly::ScalarType::Float);
|
||||
STD_TORCH_CHECK(D.is_cuda());
|
||||
STD_TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
|
||||
TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
|
||||
TORCH_CHECK(D.is_cuda());
|
||||
TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
|
||||
CHECK_SHAPE(D, dim);
|
||||
}
|
||||
|
||||
if (delta_bias_.has_value()) {
|
||||
auto delta_bias = delta_bias_.value();
|
||||
STD_TORCH_CHECK(delta_bias.scalar_type() == torch::headeronly::ScalarType::Float);
|
||||
STD_TORCH_CHECK(delta_bias.is_cuda());
|
||||
STD_TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
|
||||
TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
|
||||
TORCH_CHECK(delta_bias.is_cuda());
|
||||
TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
|
||||
CHECK_SHAPE(delta_bias, dim);
|
||||
}
|
||||
|
||||
|
||||
if (has_initial_state.has_value()) {
|
||||
auto has_initial_state_ = has_initial_state.value();
|
||||
STD_TORCH_CHECK(has_initial_state_.scalar_type() == torch::headeronly::ScalarType::Bool);
|
||||
STD_TORCH_CHECK(has_initial_state_.is_cuda());
|
||||
TORCH_CHECK(has_initial_state_.scalar_type() == at::ScalarType::Bool);
|
||||
TORCH_CHECK(has_initial_state_.is_cuda());
|
||||
CHECK_SHAPE(has_initial_state_, batch_size);
|
||||
}
|
||||
|
||||
|
||||
if (query_start_loc.has_value()) {
|
||||
auto query_start_loc_ = query_start_loc.value();
|
||||
STD_TORCH_CHECK(query_start_loc_.scalar_type() == torch::headeronly::ScalarType::Int);
|
||||
STD_TORCH_CHECK(query_start_loc_.is_cuda());
|
||||
TORCH_CHECK(query_start_loc_.scalar_type() == at::ScalarType::Int);
|
||||
TORCH_CHECK(query_start_loc_.is_cuda());
|
||||
}
|
||||
|
||||
|
||||
if (cache_indices.has_value()) {
|
||||
auto cache_indices_ = cache_indices.value();
|
||||
STD_TORCH_CHECK(cache_indices_.scalar_type() == torch::headeronly::ScalarType::Int);
|
||||
STD_TORCH_CHECK(cache_indices_.is_cuda());
|
||||
TORCH_CHECK(cache_indices_.scalar_type() == at::ScalarType::Int);
|
||||
TORCH_CHECK(cache_indices_.is_cuda());
|
||||
|
||||
// cache_indices can be either 1D (batch_size,) for non-APC mode
|
||||
// or 2D (batch_size, max_positions) for APC mode
|
||||
const bool is_apc_mode = block_idx_first_scheduled_token.has_value();
|
||||
if (is_apc_mode) {
|
||||
STD_TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode");
|
||||
STD_TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size");
|
||||
TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode");
|
||||
TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size");
|
||||
} else {
|
||||
CHECK_SHAPE(cache_indices_, batch_size);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
torch::stable::Tensor z, out_z;
|
||||
at::Tensor z, out_z;
|
||||
const bool has_z = z_.has_value();
|
||||
if (has_z) {
|
||||
z = z_.value();
|
||||
STD_TORCH_CHECK(z.scalar_type() == input_type);
|
||||
STD_TORCH_CHECK(z.is_cuda());
|
||||
STD_TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);
|
||||
TORCH_CHECK(z.scalar_type() == input_type);
|
||||
TORCH_CHECK(z.is_cuda());
|
||||
TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);
|
||||
if (varlen){
|
||||
CHECK_SHAPE(z, dim, seqlen);
|
||||
} else {
|
||||
CHECK_SHAPE(z, batch_size, dim, seqlen);
|
||||
}
|
||||
|
||||
|
||||
out_z = z;
|
||||
}
|
||||
|
||||
// Right now u has BHL layout and delta has HBL layout, and we want out to have HBL layout
|
||||
torch::stable::Tensor out = delta;
|
||||
at::Tensor out = delta;
|
||||
// ssm_states can now be either the same as input_type or float32
|
||||
auto state_type = ssm_states.scalar_type();
|
||||
STD_TORCH_CHECK(state_type == input_type || state_type == torch::headeronly::ScalarType::Float);
|
||||
STD_TORCH_CHECK(ssm_states.is_cuda());
|
||||
STD_TORCH_CHECK(ssm_states.stride(-1) == 1);
|
||||
TORCH_CHECK(state_type == input_type || state_type == at::ScalarType::Float);
|
||||
TORCH_CHECK(ssm_states.is_cuda());
|
||||
TORCH_CHECK(ssm_states.stride(-1) == 1);
|
||||
|
||||
SSMParamsBase params;
|
||||
set_ssm_params_fwd(params, batch_size, dim, seqlen, dstate, n_groups, is_variable_B, is_variable_C,
|
||||
@@ -814,8 +823,8 @@ void selective_scan_fwd(const torch::stable::Tensor &u, const torch::stable::Ten
|
||||
);
|
||||
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(u.get_device_index());
|
||||
auto stream = get_current_cuda_stream();
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(u));
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), ssm_states.scalar_type(), "selective_scan_fwd", [&] {
|
||||
selective_scan_fwd_cuda<input_t, weight_t, state_t>(params, stream);
|
||||
});
|
||||
@@ -62,9 +62,6 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
|
||||
|
||||
bool moe_permute_unpermute_supported();
|
||||
|
||||
int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
|
||||
int64_t num_experts);
|
||||
|
||||
void shuffle_rows(const torch::Tensor& input_tensor,
|
||||
const torch::Tensor& dst2src_map,
|
||||
torch::Tensor& output_tensor);
|
||||
|
||||
@@ -8,108 +8,6 @@
|
||||
// moe_permute kernels require at least CUDA 12.0
|
||||
#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
|
||||
|
||||
namespace {
|
||||
|
||||
torch::Tensor maybe_allocate_tensor(
|
||||
const std::optional<torch::Tensor>& maybe_tensor,
|
||||
at::IntArrayRef expected_sizes, torch::ScalarType dtype, c10::Device device,
|
||||
char const* name) {
|
||||
auto expected_numel = c10::multiply_integers(expected_sizes);
|
||||
if (maybe_tensor.has_value()) {
|
||||
auto tensor = maybe_tensor.value();
|
||||
TORCH_CHECK(tensor.device() == device, name, " must be on the same device");
|
||||
TORCH_CHECK(tensor.scalar_type() == dtype, name, " has incorrect dtype");
|
||||
TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
|
||||
TORCH_CHECK(tensor.numel() >= expected_numel, name,
|
||||
" is too small for the requested shape");
|
||||
auto flat_tensor = tensor.view({tensor.numel()});
|
||||
return flat_tensor.narrow(0, 0, expected_numel).view(expected_sizes);
|
||||
}
|
||||
return torch::empty(expected_sizes, torch::dtype(dtype).device(device));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
|
||||
int64_t n_expert) {
|
||||
return static_cast<int64_t>(
|
||||
CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert));
|
||||
}
|
||||
|
||||
void moe_permute_impl(
|
||||
const torch::Tensor& input, // [n_token, hidden]
|
||||
const torch::Tensor& topk_ids, // [n_token, topk]
|
||||
const torch::Tensor& token_expert_indices, // [n_token, topk]
|
||||
const std::optional<torch::Tensor>& expert_map, // [n_expert]
|
||||
int64_t n_expert, int64_t n_local_expert, int64_t topk,
|
||||
torch::Tensor& permuted_input, // [permuted_size, hidden]
|
||||
torch::Tensor& expert_first_token_offset, // [n_local_expert + 1]
|
||||
torch::Tensor& inv_permuted_idx, // [n_token, topk]
|
||||
torch::Tensor& permuted_idx, // [permute_size]
|
||||
const std::optional<torch::Tensor>& maybe_sort_workspace,
|
||||
const std::optional<torch::Tensor>& maybe_permuted_experts_id,
|
||||
const std::optional<torch::Tensor>& maybe_sorted_row_idx,
|
||||
const std::optional<torch::Tensor>& maybe_topk_ids_for_sort) {
|
||||
TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long,
|
||||
"expert_first_token_offset must be int64");
|
||||
TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
|
||||
"topk_ids must be int32");
|
||||
TORCH_CHECK(token_expert_indices.scalar_type() == at::ScalarType::Int,
|
||||
"token_expert_indices must be int32");
|
||||
TORCH_CHECK(inv_permuted_idx.scalar_type() == at::ScalarType::Int,
|
||||
"inv_permuted_idx must be int32");
|
||||
TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
|
||||
"expert_first_token_offset shape != n_local_expert+1");
|
||||
TORCH_CHECK(inv_permuted_idx.sizes() == token_expert_indices.sizes(),
|
||||
"token_expert_indices shape must be same as inv_permuted_idx");
|
||||
auto device = input.device();
|
||||
auto n_token = input.sizes()[0];
|
||||
auto n_hidden = input.sizes()[1];
|
||||
auto expanded_rows = n_token * topk;
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
|
||||
auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert);
|
||||
auto sort_workspace =
|
||||
maybe_allocate_tensor(maybe_sort_workspace, {sorter_size}, torch::kInt8,
|
||||
device, "sort_workspace");
|
||||
auto permuted_experts_id =
|
||||
maybe_allocate_tensor(maybe_permuted_experts_id, topk_ids.sizes(),
|
||||
at::ScalarType::Int, device, "permuted_experts_id");
|
||||
auto sorted_row_idx =
|
||||
maybe_allocate_tensor(maybe_sorted_row_idx, inv_permuted_idx.sizes(),
|
||||
at::ScalarType::Int, device, "sorted_row_idx");
|
||||
|
||||
CubKeyValueSorter sorter{};
|
||||
int64_t* valid_num_ptr = nullptr;
|
||||
torch::Tensor topk_ids_for_sort = topk_ids;
|
||||
|
||||
if (expert_map.has_value()) {
|
||||
const int* expert_map_ptr = get_ptr<int>(expert_map.value());
|
||||
valid_num_ptr =
|
||||
get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
|
||||
topk_ids_for_sort =
|
||||
maybe_allocate_tensor(maybe_topk_ids_for_sort, topk_ids.sizes(),
|
||||
at::ScalarType::Int, device, "topk_ids_for_sort");
|
||||
topk_ids_for_sort.copy_(topk_ids);
|
||||
preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
|
||||
expert_map_ptr, n_expert, stream);
|
||||
}
|
||||
|
||||
sortAndScanExpert(
|
||||
get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
|
||||
get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
|
||||
get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
|
||||
n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
|
||||
|
||||
MOE_DISPATCH(input.scalar_type(), [&] {
|
||||
expandInputRowsKernelLauncher<scalar_t>(
|
||||
get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
|
||||
get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
|
||||
get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
|
||||
n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
|
||||
});
|
||||
}
|
||||
|
||||
void moe_permute(
|
||||
const torch::Tensor& input, // [n_token, hidden]
|
||||
const torch::Tensor& topk_ids, // [n_token, topk]
|
||||
@@ -120,26 +18,65 @@ void moe_permute(
|
||||
torch::Tensor& expert_first_token_offset, // [n_local_expert + 1]
|
||||
torch::Tensor& inv_permuted_idx, // [n_token, topk]
|
||||
torch::Tensor& permuted_idx) { // [permute_size]
|
||||
moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
|
||||
n_local_expert, topk, permuted_input,
|
||||
expert_first_token_offset, inv_permuted_idx, permuted_idx,
|
||||
std::nullopt, std::nullopt, std::nullopt, std::nullopt);
|
||||
}
|
||||
TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long,
|
||||
"expert_first_token_offset must be int64");
|
||||
TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
|
||||
"topk_ids must be int32");
|
||||
TORCH_CHECK(token_expert_indices.scalar_type() == at::ScalarType::Int,
|
||||
"token_expert_indices must be int32");
|
||||
TORCH_CHECK(inv_permuted_idx.scalar_type() == at::ScalarType::Int,
|
||||
"inv_permuted_idx must be int32");
|
||||
TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
|
||||
"expert_first_token_offset shape != n_local_expert+1")
|
||||
TORCH_CHECK(inv_permuted_idx.sizes() == token_expert_indices.sizes(),
|
||||
"token_expert_indices shape must be same as inv_permuted_idx");
|
||||
auto n_token = input.sizes()[0];
|
||||
auto n_hidden = input.sizes()[1];
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
const long sorter_size =
|
||||
CubKeyValueSorter::getWorkspaceSize(n_token * topk, n_expert);
|
||||
auto sort_workspace = torch::empty(
|
||||
{sorter_size},
|
||||
torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false));
|
||||
torch::Tensor topk_ids_for_sort = topk_ids;
|
||||
auto permuted_experts_id = torch::empty_like(topk_ids);
|
||||
auto sorted_row_idx = torch::empty_like(inv_permuted_idx);
|
||||
|
||||
void moe_permute_with_scratch(
|
||||
const torch::Tensor& input, const torch::Tensor& topk_ids,
|
||||
const torch::Tensor& token_expert_indices,
|
||||
const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
|
||||
int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
|
||||
torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
|
||||
torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
|
||||
torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
|
||||
torch::Tensor& topk_ids_for_sort) {
|
||||
moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
|
||||
n_local_expert, topk, permuted_input,
|
||||
expert_first_token_offset, inv_permuted_idx, permuted_idx,
|
||||
sort_workspace, permuted_experts_id, sorted_row_idx,
|
||||
topk_ids_for_sort);
|
||||
CubKeyValueSorter sorter{};
|
||||
int64_t* valid_num_ptr = nullptr;
|
||||
// pre-process kernel for expert-parallelism:
|
||||
// no local expert id plus "n_expert" offset for priority to local expert
|
||||
// map local expert id [n, .., n+n_local_expert-1] to [0, n_local_expert -1]
|
||||
// For example, 4 expert with ep_size=2. ep_rank=1 owns global expert id
|
||||
// [2,3] with expert_map[-1, -1, 0, 1], preprocess_topk_id process topk_ids
|
||||
// and map global expert id [2, 3] to local_expert id [0, 1] and map global
|
||||
// expert id [0, 1] ( not in ep rank=1) to [4, 5] by plus n_expert. This map
|
||||
// operation is to make local expert high priority in following sort topk_ids
|
||||
// and scan local expert_first_token_offset for each ep rank for next group
|
||||
// gemm.
|
||||
if (expert_map.has_value()) {
|
||||
const int* expert_map_ptr = get_ptr<int>(expert_map.value());
|
||||
valid_num_ptr =
|
||||
get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
|
||||
topk_ids_for_sort = topk_ids.clone();
|
||||
preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
|
||||
expert_map_ptr, n_expert, stream);
|
||||
}
|
||||
// expert sort topk expert id and scan expert id get expert_first_token_offset
|
||||
sortAndScanExpert(
|
||||
get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
|
||||
get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
|
||||
get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
|
||||
n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
|
||||
|
||||
// dispatch expandInputRowsKernelLauncher
|
||||
MOE_DISPATCH(input.scalar_type(), [&] {
|
||||
expandInputRowsKernelLauncher<scalar_t>(
|
||||
get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
|
||||
get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
|
||||
get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
|
||||
n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
|
||||
});
|
||||
}
|
||||
|
||||
void moe_unpermute(
|
||||
@@ -232,12 +169,6 @@ void shuffle_rows(const torch::Tensor& input_tensor,
|
||||
|
||||
#else
|
||||
|
||||
int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
|
||||
int64_t n_expert) {
|
||||
TORCH_CHECK(
|
||||
false, "moe_permute_sort_workspace_size is not supported on CUDA < 12.0");
|
||||
}
|
||||
|
||||
void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
|
||||
const torch::Tensor& token_expert_indices,
|
||||
const std::optional<torch::Tensor>& expert_map,
|
||||
@@ -248,19 +179,6 @@ void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
|
||||
TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0");
|
||||
}
|
||||
|
||||
void moe_permute_with_scratch(
|
||||
const torch::Tensor& input, const torch::Tensor& topk_ids,
|
||||
const torch::Tensor& token_expert_indices,
|
||||
const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
|
||||
int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
|
||||
torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
|
||||
torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
|
||||
torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
|
||||
torch::Tensor& topk_ids_for_sort) {
|
||||
TORCH_CHECK(false,
|
||||
"moe_permute_with_scratch is not supported on CUDA < 12.0");
|
||||
}
|
||||
|
||||
void moe_unpermute(
|
||||
const torch::Tensor& permuted_hidden_states,
|
||||
const torch::Tensor& topk_weights, const torch::Tensor& inv_permuted_idx,
|
||||
@@ -281,6 +199,5 @@ bool moe_permute_unpermute_supported() {
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("moe_permute", &moe_permute);
|
||||
m.impl("moe_permute_with_scratch", &moe_permute_with_scratch);
|
||||
m.impl("moe_unpermute", &moe_unpermute);
|
||||
}
|
||||
@@ -100,26 +100,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
|
||||
"expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
|
||||
"permuted_idx)->()");
|
||||
|
||||
m.def(
|
||||
"moe_permute_with_scratch(Tensor input, Tensor topk_ids,"
|
||||
"Tensor token_expert_indices, Tensor? expert_map, int n_expert,"
|
||||
"int n_local_expert,"
|
||||
"int topk, Tensor! permuted_input, Tensor! "
|
||||
"expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
|
||||
"permuted_idx, Tensor! sort_workspace, Tensor! permuted_experts_id, "
|
||||
"Tensor! sorted_row_idx, Tensor! topk_ids_for_sort)->()");
|
||||
|
||||
m.def(
|
||||
"moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights,"
|
||||
"Tensor inv_permuted_idx, Tensor? expert_first_token_offset, "
|
||||
"int topk, Tensor! hidden_states)->()");
|
||||
|
||||
m.def("moe_permute_unpermute_supported() -> bool");
|
||||
m.def(
|
||||
"moe_permute_sort_workspace_size(int num_expanded_rows, int n_expert) -> "
|
||||
"int");
|
||||
m.impl("moe_permute_unpermute_supported", &moe_permute_unpermute_supported);
|
||||
m.impl("moe_permute_sort_workspace_size", &moe_permute_sort_workspace_size);
|
||||
|
||||
// Row shuffle for MoE
|
||||
m.def(
|
||||
|
||||
+43
@@ -54,6 +54,13 @@ void paged_attention_v2(
|
||||
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
|
||||
const int64_t blocksparse_head_sliding_step);
|
||||
|
||||
void merge_attn_states(
|
||||
torch::Tensor& output, std::optional<torch::Tensor> output_lse,
|
||||
const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
|
||||
const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
|
||||
const std::optional<int64_t> prefill_tokens_with_context,
|
||||
const std::optional<torch::Tensor>& output_scale = std::nullopt);
|
||||
|
||||
// rms_norm and fused_add_rms_norm declarations also exist in
|
||||
// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
|
||||
// because the CPU build still uses these torch::Tensor declarations.
|
||||
@@ -69,6 +76,26 @@ torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
|
||||
torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
|
||||
int64_t cache_block_size);
|
||||
|
||||
void apply_repetition_penalties_(torch::Tensor& logits,
|
||||
const torch::Tensor& prompt_mask,
|
||||
const torch::Tensor& output_mask,
|
||||
const torch::Tensor& repetition_penalties);
|
||||
|
||||
void top_k_per_row_prefill(const torch::Tensor& logits,
|
||||
const torch::Tensor& rowStarts,
|
||||
const torch::Tensor& rowEnds, torch::Tensor& indices,
|
||||
int64_t numRows, int64_t stride0, int64_t stride1,
|
||||
int64_t topK);
|
||||
|
||||
void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
|
||||
const torch::Tensor& seqLens, torch::Tensor& indices,
|
||||
int64_t numRows, int64_t stride0, int64_t stride1,
|
||||
int64_t topK);
|
||||
|
||||
void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
torch::Tensor& output, torch::Tensor& workspace, int64_t k,
|
||||
int64_t max_seq_len);
|
||||
|
||||
void silu_and_mul_per_block_quant(torch::Tensor& out,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor& scales, int64_t group_size,
|
||||
@@ -123,6 +150,22 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
|
||||
torch::Tensor& scales,
|
||||
std::optional<torch::Tensor> const& azp);
|
||||
|
||||
void selective_scan_fwd(
|
||||
const torch::Tensor& u, const torch::Tensor& delta, const torch::Tensor& A,
|
||||
const torch::Tensor& B, const torch::Tensor& C,
|
||||
const std::optional<torch::Tensor>& D_,
|
||||
const std::optional<torch::Tensor>& z_,
|
||||
const std::optional<torch::Tensor>& delta_bias_, bool delta_softplus,
|
||||
const std::optional<torch::Tensor>& query_start_loc,
|
||||
const std::optional<torch::Tensor>& cache_indices,
|
||||
const std::optional<torch::Tensor>& has_initial_state,
|
||||
const torch::Tensor& ssm_states, int64_t null_block_id, int64_t block_size,
|
||||
const std::optional<torch::Tensor>& block_idx_first_scheduled_token,
|
||||
const std::optional<torch::Tensor>& block_idx_last_scheduled_token,
|
||||
const std::optional<torch::Tensor>& initial_state_idx,
|
||||
const std::optional<torch::Tensor>& cu_chunk_seqlen,
|
||||
const std::optional<torch::Tensor>& last_chunk_indices);
|
||||
|
||||
torch::Tensor dynamic_4bit_int_moe_cpu(
|
||||
torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
|
||||
torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
|
||||
|
||||
@@ -126,10 +126,10 @@ struct RadixRowState {
|
||||
// ============================================================================
|
||||
|
||||
struct PersistentTopKParams {
|
||||
const float* __restrict__ input; // [num_rows, stride]
|
||||
int32_t* __restrict__ output; // [num_rows, top_k]
|
||||
const int32_t* __restrict__ lengths; // [num_rows]
|
||||
RadixRowState* row_states; // large path: per-group state
|
||||
const float* __restrict__ input; // [num_rows, stride]
|
||||
int32_t* __restrict__ output; // [num_rows, top_k]
|
||||
int32_t* __restrict__ lengths; // [num_rows]
|
||||
RadixRowState* row_states; // large path: per-group state
|
||||
uint32_t num_rows;
|
||||
uint32_t stride;
|
||||
uint32_t top_k; // actual k value for output stride
|
||||
@@ -1269,11 +1269,9 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
|
||||
}
|
||||
|
||||
template <typename DType, typename IdType, uint32_t MAX_K = 2048>
|
||||
cudaError_t FilteredTopKRaggedTransform(const DType* input,
|
||||
IdType* output_indices,
|
||||
const IdType* lengths,
|
||||
uint32_t num_rows, uint32_t top_k_val,
|
||||
uint32_t max_len,
|
||||
cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
|
||||
IdType* lengths, uint32_t num_rows,
|
||||
uint32_t top_k_val, uint32_t max_len,
|
||||
cudaStream_t stream = 0) {
|
||||
constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC;
|
||||
constexpr int MAX_VEC = 16 / sizeof(DType);
|
||||
|
||||
@@ -400,12 +400,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
|
||||
"Turing only support FP16 or INT8 activation.");
|
||||
}
|
||||
if (a_type == vllm::kFE4M3fn) {
|
||||
TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
|
||||
"FP8 only support Ada Lovelace or newer GPUs.");
|
||||
TORCH_CHECK(
|
||||
major_capability * 10 + minor_capability == 89 ||
|
||||
major_capability == 12,
|
||||
"Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than "
|
||||
major_capability * 10 + minor_capability == 120,
|
||||
"Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than "
|
||||
"Marlin W4A16 on other devices).");
|
||||
}
|
||||
|
||||
|
||||
@@ -1277,12 +1277,6 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
|
||||
else
|
||||
WVSPLIT_TILE_CFG(64, 16, sYT, 4)
|
||||
break;
|
||||
case 5:
|
||||
if (use_wave32)
|
||||
WVSPLIT_TILE_CFG(32, 16, sYT, 5)
|
||||
else
|
||||
WVSPLIT_TILE_CFG(64, 16, sYT, 5)
|
||||
break;
|
||||
default:
|
||||
throw std::runtime_error(
|
||||
"Unsupported N value: " + std::to_string(M_in) + "," +
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
#include "../cuda_compat.h"
|
||||
#include "cuda_compat.h"
|
||||
#include "dispatch_utils.h"
|
||||
#include "torch_utils.h"
|
||||
|
||||
#include <torch/cuda.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#include <cub/cub.cuh>
|
||||
@@ -616,14 +618,14 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
|
||||
} // namespace vllm
|
||||
|
||||
void apply_repetition_penalties_(
|
||||
torch::stable::Tensor& logits, // [num_seqs, vocab_size], in-place
|
||||
const torch::stable::Tensor& prompt_mask, // [num_seqs, vocab_size]
|
||||
const torch::stable::Tensor& output_mask, // [num_seqs, vocab_size]
|
||||
const torch::stable::Tensor& repetition_penalties) { // [num_seqs]
|
||||
STD_TORCH_CHECK(logits.is_contiguous());
|
||||
STD_TORCH_CHECK(prompt_mask.is_contiguous());
|
||||
STD_TORCH_CHECK(output_mask.is_contiguous());
|
||||
STD_TORCH_CHECK(repetition_penalties.is_contiguous());
|
||||
torch::Tensor& logits, // [num_seqs, vocab_size], in-place
|
||||
const torch::Tensor& prompt_mask, // [num_seqs, vocab_size]
|
||||
const torch::Tensor& output_mask, // [num_seqs, vocab_size]
|
||||
const torch::Tensor& repetition_penalties) { // [num_seqs]
|
||||
TORCH_CHECK(logits.is_contiguous());
|
||||
TORCH_CHECK(prompt_mask.is_contiguous());
|
||||
TORCH_CHECK(output_mask.is_contiguous());
|
||||
TORCH_CHECK(repetition_penalties.is_contiguous());
|
||||
|
||||
int vocab_size = logits.size(-1);
|
||||
int num_seqs = logits.size(0);
|
||||
@@ -633,7 +635,7 @@ void apply_repetition_penalties_(
|
||||
// Get number of SMs on the current device
|
||||
int sms = 0;
|
||||
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount,
|
||||
logits.get_device_index());
|
||||
logits.get_device());
|
||||
|
||||
// Compute tile_num and tile_size
|
||||
int tile_num =
|
||||
@@ -643,29 +645,27 @@ void apply_repetition_penalties_(
|
||||
// Each block handles one sequence and a tile of vocab
|
||||
dim3 grid(num_seqs, tile_num);
|
||||
dim3 block(std::min(tile_size, 1024));
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
logits.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream();
|
||||
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(logits));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
logits.scalar_type(), "apply_repetition_penalties_kernel", [&] {
|
||||
vllm::apply_repetition_penalties_kernel<scalar_t>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
logits.mutable_data_ptr<scalar_t>(),
|
||||
prompt_mask.const_data_ptr<bool>(),
|
||||
output_mask.const_data_ptr<bool>(),
|
||||
repetition_penalties.const_data_ptr<scalar_t>(), num_seqs,
|
||||
vocab_size, tile_size);
|
||||
logits.data_ptr<scalar_t>(), prompt_mask.data_ptr<bool>(),
|
||||
output_mask.data_ptr<bool>(),
|
||||
repetition_penalties.data_ptr<scalar_t>(), num_seqs, vocab_size,
|
||||
tile_size);
|
||||
});
|
||||
}
|
||||
|
||||
void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
|
||||
const torch::stable::Tensor& seqLens,
|
||||
torch::stable::Tensor& indices, int64_t numRows,
|
||||
int64_t stride0, int64_t stride1, int64_t topK) {
|
||||
void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
|
||||
const torch::Tensor& seqLens, torch::Tensor& indices,
|
||||
int64_t numRows, int64_t stride0, int64_t stride1,
|
||||
int64_t topK) {
|
||||
constexpr int kSortingAlgorithmThreshold = 12288;
|
||||
constexpr int kSplitWorkThreshold = 200 * 1000;
|
||||
constexpr int kNumThreadsPerBlock = 512;
|
||||
const cudaStream_t stream = get_current_cuda_stream();
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
const auto numColumns = logits.size(1);
|
||||
|
||||
// True if seqLens is 2D (B, next_n): each logit row has its own pre-computed
|
||||
@@ -677,76 +677,73 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
|
||||
// Use insertion sort
|
||||
vllm::topKPerRowDecode<kNumThreadsPerBlock, false>
|
||||
<<<numRows, kNumThreadsPerBlock, topK * sizeof(int32_t), stream>>>(
|
||||
logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
|
||||
indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
|
||||
logits.data_ptr<float>(), seqLens.data_ptr<int>(),
|
||||
indices.data_ptr<int>(), static_cast<int>(stride0),
|
||||
static_cast<int>(stride1), static_cast<int>(topK),
|
||||
static_cast<int>(next_n), seqLensIs2D);
|
||||
} else if (numColumns < kSplitWorkThreshold) {
|
||||
// From this threshold, use radix sort instead
|
||||
vllm::topKPerRowDecode<kNumThreadsPerBlock, true>
|
||||
<<<numRows, kNumThreadsPerBlock, topK * sizeof(int32_t), stream>>>(
|
||||
logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
|
||||
indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
|
||||
logits.data_ptr<float>(), seqLens.data_ptr<int>(),
|
||||
indices.data_ptr<int>(), static_cast<int>(stride0),
|
||||
static_cast<int>(stride1), static_cast<int>(topK),
|
||||
static_cast<int>(next_n), seqLensIs2D);
|
||||
} else {
|
||||
// Long sequences are run in two steps
|
||||
constexpr auto multipleBlocksPerRowConfig = 10;
|
||||
|
||||
const auto outIndicesAux = torch::stable::empty(
|
||||
{numRows, multipleBlocksPerRowConfig, topK},
|
||||
torch::headeronly::ScalarType::Int, std::nullopt, logits.device());
|
||||
const auto outLogitsAux = torch::stable::empty(
|
||||
{numRows, multipleBlocksPerRowConfig, topK},
|
||||
torch::headeronly::ScalarType::Float, std::nullopt, logits.device());
|
||||
const auto outIndicesAux =
|
||||
torch::empty({numRows, multipleBlocksPerRowConfig, topK},
|
||||
torch::dtype(torch::kInt32).device(logits.device()));
|
||||
const auto outLogitsAux =
|
||||
torch::empty({numRows, multipleBlocksPerRowConfig, topK},
|
||||
torch::dtype(torch::kFloat).device(logits.device()));
|
||||
|
||||
vllm::topKPerRowDecode<kNumThreadsPerBlock, true, true>
|
||||
<<<dim3(numRows, multipleBlocksPerRowConfig), kNumThreadsPerBlock,
|
||||
2 * topK * sizeof(int32_t), stream>>>(
|
||||
logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
|
||||
outIndicesAux.mutable_data_ptr<int>(), static_cast<int>(stride0),
|
||||
logits.data_ptr<float>(), seqLens.data_ptr<int>(),
|
||||
outIndicesAux.data_ptr<int>(), static_cast<int>(stride0),
|
||||
static_cast<int>(stride1), static_cast<int>(topK),
|
||||
static_cast<int>(next_n), seqLensIs2D,
|
||||
outLogitsAux.mutable_data_ptr<float>());
|
||||
outLogitsAux.data_ptr<float>());
|
||||
|
||||
constexpr int kNumThreadsPerBlockMerge = 1024;
|
||||
vllm::topKPerRowDecode<kNumThreadsPerBlockMerge, true, false, true>
|
||||
<<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t), stream>>>(
|
||||
outLogitsAux.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
|
||||
indices.mutable_data_ptr<int>(), multipleBlocksPerRowConfig * topK,
|
||||
1, static_cast<int>(topK), static_cast<int>(next_n), seqLensIs2D,
|
||||
nullptr, multipleBlocksPerRowConfig,
|
||||
outIndicesAux.const_data_ptr<int>());
|
||||
outLogitsAux.data_ptr<float>(), seqLens.data_ptr<int>(),
|
||||
indices.data_ptr<int>(), multipleBlocksPerRowConfig * topK, 1,
|
||||
static_cast<int>(topK), static_cast<int>(next_n), seqLensIs2D,
|
||||
nullptr, multipleBlocksPerRowConfig, outIndicesAux.data_ptr<int>());
|
||||
}
|
||||
}
|
||||
|
||||
void top_k_per_row_prefill(const torch::stable::Tensor& logits,
|
||||
const torch::stable::Tensor& rowStarts,
|
||||
const torch::stable::Tensor& rowEnds,
|
||||
torch::stable::Tensor& indices, int64_t numRows,
|
||||
int64_t stride0, int64_t stride1, int64_t topK) {
|
||||
void top_k_per_row_prefill(const torch::Tensor& logits,
|
||||
const torch::Tensor& rowStarts,
|
||||
const torch::Tensor& rowEnds, torch::Tensor& indices,
|
||||
int64_t numRows, int64_t stride0, int64_t stride1,
|
||||
int64_t topK) {
|
||||
constexpr int kSortingAlgorithmThreshold = 12288;
|
||||
constexpr int kNumThreadsPerBlock = 512;
|
||||
const cudaStream_t stream = get_current_cuda_stream();
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
int numInsertionBlocks =
|
||||
std::min(static_cast<int>(numRows), kSortingAlgorithmThreshold);
|
||||
vllm::topKPerRowPrefill<kNumThreadsPerBlock, false>
|
||||
<<<numInsertionBlocks, kNumThreadsPerBlock, topK * sizeof(int32_t),
|
||||
stream>>>(logits.const_data_ptr<float>(),
|
||||
rowStarts.const_data_ptr<int>(),
|
||||
rowEnds.const_data_ptr<int>(),
|
||||
indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
|
||||
static_cast<int>(stride1), static_cast<int>(topK), 0);
|
||||
stream>>>(logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
|
||||
rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
|
||||
static_cast<int>(stride0), static_cast<int>(stride1),
|
||||
static_cast<int>(topK), 0);
|
||||
|
||||
if (numRows > kSortingAlgorithmThreshold) {
|
||||
int numRadixBlocks = numRows - kSortingAlgorithmThreshold;
|
||||
vllm::topKPerRowPrefill<kNumThreadsPerBlock, true>
|
||||
<<<numRadixBlocks, kNumThreadsPerBlock, topK * sizeof(int32_t),
|
||||
stream>>>(
|
||||
logits.const_data_ptr<float>(), rowStarts.const_data_ptr<int>(),
|
||||
rowEnds.const_data_ptr<int>(), indices.mutable_data_ptr<int>(),
|
||||
static_cast<int>(stride0), static_cast<int>(stride1),
|
||||
static_cast<int>(topK), kSortingAlgorithmThreshold);
|
||||
stream>>>(logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
|
||||
rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
|
||||
static_cast<int>(stride0), static_cast<int>(stride1),
|
||||
static_cast<int>(topK), kSortingAlgorithmThreshold);
|
||||
}
|
||||
}
|
||||
@@ -1,51 +1,49 @@
|
||||
// Persistent TopK kernel for DeepSeek V3 sparse attention indexer.
|
||||
// See persistent_topk.cuh for kernel implementation.
|
||||
|
||||
#include <torch/all.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <algorithm>
|
||||
|
||||
#include "torch_utils.h"
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#include "../persistent_topk.cuh"
|
||||
#include "persistent_topk.cuh"
|
||||
#endif
|
||||
|
||||
namespace {
|
||||
|
||||
#ifndef USE_ROCM
|
||||
template <int TopK>
|
||||
void launch_persistent_topk(const torch::stable::Tensor& logits,
|
||||
const torch::stable::Tensor& lengths,
|
||||
torch::stable::Tensor& output,
|
||||
torch::stable::Tensor& workspace,
|
||||
int64_t max_seq_len) {
|
||||
void launch_persistent_topk(const torch::Tensor& logits,
|
||||
const torch::Tensor& lengths, torch::Tensor& output,
|
||||
torch::Tensor& workspace, int64_t max_seq_len) {
|
||||
namespace P = vllm::persistent;
|
||||
|
||||
const int64_t num_rows = logits.size(0);
|
||||
const int64_t stride = logits.stride(0);
|
||||
const cudaStream_t stream = get_current_cuda_stream();
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
static int num_sms = 0;
|
||||
static int max_smem_per_block = 0;
|
||||
if (num_sms == 0) {
|
||||
const cudaDeviceProp* device_prop = get_device_prop();
|
||||
num_sms = device_prop->multiProcessorCount;
|
||||
max_smem_per_block = device_prop->sharedMemPerBlockOptin;
|
||||
int device;
|
||||
cudaGetDevice(&device);
|
||||
cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device);
|
||||
cudaDeviceGetAttribute(&max_smem_per_block,
|
||||
cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
|
||||
}
|
||||
|
||||
if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
|
||||
cudaError_t status =
|
||||
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
|
||||
logits.const_data_ptr<float>(), output.mutable_data_ptr<int32_t>(),
|
||||
lengths.const_data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
|
||||
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
|
||||
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
|
||||
static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
|
||||
STD_TORCH_CHECK(status == cudaSuccess,
|
||||
"FilteredTopK failed: ", cudaGetErrorString(status));
|
||||
TORCH_CHECK(status == cudaSuccess,
|
||||
"FilteredTopK failed: ", cudaGetErrorString(status));
|
||||
} else {
|
||||
STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
|
||||
STD_TORCH_CHECK(
|
||||
workspace.scalar_type() == torch::headeronly::ScalarType::Byte,
|
||||
"workspace must be uint8");
|
||||
TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
|
||||
TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
|
||||
|
||||
int effective_max_smem;
|
||||
if (num_rows <= 4) {
|
||||
@@ -101,9 +99,9 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
|
||||
&occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
|
||||
smem_size);
|
||||
}
|
||||
STD_TORCH_CHECK(occ_err == cudaSuccess,
|
||||
"persistent_topk occupancy query failed: ",
|
||||
cudaGetErrorString(occ_err));
|
||||
TORCH_CHECK(occ_err == cudaSuccess,
|
||||
"persistent_topk occupancy query failed: ",
|
||||
cudaGetErrorString(occ_err));
|
||||
if (occupancy < 1) occupancy = 1;
|
||||
|
||||
// The cooperative spin-wait barrier only runs when at least one row hits
|
||||
@@ -133,29 +131,27 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
|
||||
// If the cooperative launch wouldn't fit, fall back to FilteredTopK
|
||||
// instead of deadlocking. Only relevant when needs_cooperative.
|
||||
if (needs_cooperative && total_ctas > hw_resident_cap) {
|
||||
STD_TORCH_CHECK(
|
||||
max_smem_per_block >= 128 * 1024,
|
||||
"persistent_topk would oversubscribe and the FilteredTopK "
|
||||
"fallback requires >=128KB smem per block (have ",
|
||||
max_smem_per_block, "). total_ctas=", total_ctas,
|
||||
" > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
|
||||
", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
|
||||
", smem=", smem_size, ").");
|
||||
TORCH_CHECK(max_smem_per_block >= 128 * 1024,
|
||||
"persistent_topk would oversubscribe and the FilteredTopK "
|
||||
"fallback requires >=128KB smem per block (have ",
|
||||
max_smem_per_block, "). total_ctas=", total_ctas,
|
||||
" > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
|
||||
", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
|
||||
", smem=", smem_size, ").");
|
||||
cudaError_t status =
|
||||
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
|
||||
logits.const_data_ptr<float>(),
|
||||
output.mutable_data_ptr<int32_t>(),
|
||||
lengths.const_data_ptr<int32_t>(),
|
||||
static_cast<uint32_t>(num_rows), static_cast<uint32_t>(TopK),
|
||||
static_cast<uint32_t>(stride), stream);
|
||||
STD_TORCH_CHECK(status == cudaSuccess, "FilteredTopK fallback failed: ",
|
||||
cudaGetErrorString(status));
|
||||
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
|
||||
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
|
||||
static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
|
||||
stream);
|
||||
TORCH_CHECK(status == cudaSuccess,
|
||||
"FilteredTopK fallback failed: ", cudaGetErrorString(status));
|
||||
return;
|
||||
}
|
||||
|
||||
size_t state_bytes = num_groups * sizeof(P::RadixRowState);
|
||||
STD_TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
|
||||
"workspace too small, need ", state_bytes, " bytes");
|
||||
TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
|
||||
"workspace too small, need ", state_bytes, " bytes");
|
||||
|
||||
// Zero the per-group RadixRowState region before launch.
|
||||
//
|
||||
@@ -183,22 +179,22 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
|
||||
// first red_release. cudaMemsetAsync is stream-ordered: the zero
|
||||
// is globally visible before any CTA runs.
|
||||
{
|
||||
cudaError_t mz_err = cudaMemsetAsync(
|
||||
workspace.mutable_data_ptr<uint8_t>(), 0, state_bytes, stream);
|
||||
STD_TORCH_CHECK(mz_err == cudaSuccess,
|
||||
"row_states memset failed: ", cudaGetErrorString(mz_err));
|
||||
cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr<uint8_t>(), 0,
|
||||
state_bytes, stream);
|
||||
TORCH_CHECK(mz_err == cudaSuccess,
|
||||
"row_states memset failed: ", cudaGetErrorString(mz_err));
|
||||
}
|
||||
|
||||
P::PersistentTopKParams params;
|
||||
params.input = logits.const_data_ptr<float>();
|
||||
params.output = output.mutable_data_ptr<int32_t>();
|
||||
params.lengths = lengths.const_data_ptr<int32_t>();
|
||||
params.input = logits.data_ptr<float>();
|
||||
params.output = output.data_ptr<int32_t>();
|
||||
params.lengths = lengths.data_ptr<int32_t>();
|
||||
params.num_rows = static_cast<uint32_t>(num_rows);
|
||||
params.stride = static_cast<uint32_t>(stride);
|
||||
params.top_k = static_cast<uint32_t>(TopK);
|
||||
params.chunk_size = chunk_size;
|
||||
params.row_states = reinterpret_cast<P::RadixRowState*>(
|
||||
workspace.mutable_data_ptr<uint8_t>());
|
||||
params.row_states =
|
||||
reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
|
||||
params.ctas_per_group = ctas_per_group;
|
||||
params.max_seq_len = static_cast<uint32_t>(max_seq_len);
|
||||
|
||||
@@ -207,8 +203,8 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
|
||||
auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>; \
|
||||
cudaError_t err = cudaFuncSetAttribute( \
|
||||
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \
|
||||
STD_TORCH_CHECK(err == cudaSuccess, \
|
||||
"Failed to set smem: ", cudaGetErrorString(err)); \
|
||||
TORCH_CHECK(err == cudaSuccess, \
|
||||
"Failed to set smem: ", cudaGetErrorString(err)); \
|
||||
kernel<<<total_ctas, P::kThreadsPerBlock, smem_size, stream>>>(params); \
|
||||
} while (0)
|
||||
|
||||
@@ -223,42 +219,37 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
|
||||
}
|
||||
|
||||
cudaError_t err = cudaGetLastError();
|
||||
STD_TORCH_CHECK(err == cudaSuccess,
|
||||
"persistent_topk failed: ", cudaGetErrorString(err));
|
||||
TORCH_CHECK(err == cudaSuccess,
|
||||
"persistent_topk failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
#endif
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
void persistent_topk(const torch::stable::Tensor& logits,
|
||||
const torch::stable::Tensor& lengths,
|
||||
torch::stable::Tensor& output,
|
||||
torch::stable::Tensor& workspace, int64_t k,
|
||||
void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
torch::Tensor& output, torch::Tensor& workspace, int64_t k,
|
||||
int64_t max_seq_len) {
|
||||
#ifndef USE_ROCM
|
||||
STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
|
||||
STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
|
||||
STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
|
||||
STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float,
|
||||
"Only float32 supported");
|
||||
STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"lengths must be int32");
|
||||
STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"output must be int32");
|
||||
STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
|
||||
STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
|
||||
"lengths must be 1D or 2D");
|
||||
STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
|
||||
STD_TORCH_CHECK(output.dim() == 2, "output must be 2D");
|
||||
TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
|
||||
TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
|
||||
TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
|
||||
TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
|
||||
TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
|
||||
TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
|
||||
TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
|
||||
TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
|
||||
"lengths must be 1D or 2D");
|
||||
TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
|
||||
TORCH_CHECK(output.dim() == 2, "output must be 2D");
|
||||
|
||||
const int64_t num_rows = logits.size(0);
|
||||
const int64_t stride = logits.stride(0);
|
||||
|
||||
STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
|
||||
STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
|
||||
"output size mismatch");
|
||||
STD_TORCH_CHECK(
|
||||
k == 512 || k == 1024 || k == 2048,
|
||||
"persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
|
||||
TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
|
||||
TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
|
||||
"output size mismatch");
|
||||
TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
|
||||
"persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
|
||||
|
||||
if (k == 512) {
|
||||
launch_persistent_topk<512>(logits, lengths, output, workspace,
|
||||
@@ -271,6 +262,6 @@ void persistent_topk(const torch::stable::Tensor& logits,
|
||||
max_seq_len);
|
||||
}
|
||||
#else
|
||||
STD_TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
|
||||
TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
|
||||
#endif
|
||||
}
|
||||
@@ -62,6 +62,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
" int blocksparse_head_sliding_step) -> ()");
|
||||
ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);
|
||||
|
||||
// Merge attn states
|
||||
// Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
|
||||
// can be used to combine partial attention results (in the split-KV case)
|
||||
ops.def(
|
||||
"merge_attn_states("
|
||||
" Tensor! output,"
|
||||
" Tensor!? output_lse,"
|
||||
" Tensor prefix_output,"
|
||||
" Tensor prefix_lse,"
|
||||
" Tensor suffix_output,"
|
||||
" Tensor suffix_lse,"
|
||||
" int!? prefill_tokens_with_context,"
|
||||
" Tensor? output_scale=None) -> ()");
|
||||
ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
|
||||
|
||||
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
|
||||
ops.def(
|
||||
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
||||
@@ -90,6 +105,31 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
|
||||
&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
|
||||
|
||||
// Apply repetition penalties to logits in-place
|
||||
ops.def(
|
||||
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
|
||||
"Tensor output_mask, Tensor repetition_penalties) -> ()");
|
||||
ops.impl("apply_repetition_penalties_", torch::kCUDA,
|
||||
&apply_repetition_penalties_);
|
||||
|
||||
// Optimized top-k per row operation
|
||||
ops.def(
|
||||
"top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
|
||||
"Tensor! indices, int numRows, int stride0, "
|
||||
"int stride1, int topK) -> ()");
|
||||
ops.impl("top_k_per_row_prefill", torch::kCUDA, &top_k_per_row_prefill);
|
||||
|
||||
ops.def(
|
||||
"top_k_per_row_decode(Tensor logits, int next_n, "
|
||||
"Tensor seq_lens, Tensor! indices, "
|
||||
"int numRows, int stride0, int stride1, int topK) -> ()");
|
||||
ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode);
|
||||
|
||||
ops.def(
|
||||
"persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
|
||||
"Tensor workspace, int k, int max_seq_len) -> ()");
|
||||
ops.impl("persistent_topk", torch::kCUDA, &persistent_topk);
|
||||
|
||||
// Quantization ops
|
||||
#ifndef USE_ROCM
|
||||
|
||||
@@ -190,6 +230,25 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
|
||||
#endif
|
||||
|
||||
// Mamba selective scan kernel
|
||||
ops.def(
|
||||
"selective_scan_fwd(Tensor! u, Tensor! delta,"
|
||||
"Tensor! A, Tensor! B, Tensor! C,"
|
||||
"Tensor? D_, Tensor!? z_, Tensor? delta_bias_,"
|
||||
"bool delta_softplus,"
|
||||
"Tensor? query_start_loc,"
|
||||
"Tensor? cache_indices,"
|
||||
"Tensor? has_initial_state,"
|
||||
"Tensor! ssm_states,"
|
||||
"int null_block_id,"
|
||||
"int block_size,"
|
||||
"Tensor? block_idx_first_scheduled_token,"
|
||||
"Tensor? block_idx_last_scheduled_token,"
|
||||
"Tensor? initial_state_idx,"
|
||||
"Tensor? cu_chunk_seqlen,"
|
||||
"Tensor? last_chunk_indices) -> ()");
|
||||
ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
ops.def(
|
||||
"minimax_allreduce_rms("
|
||||
|
||||
@@ -868,7 +868,6 @@ FROM vllm-base AS test
|
||||
ADD . /vllm-workspace/
|
||||
|
||||
ARG PYTHON_VERSION
|
||||
ARG TARGETPLATFORM
|
||||
|
||||
ARG PIP_INDEX_URL UV_INDEX_URL
|
||||
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
|
||||
@@ -908,11 +907,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
|
||||
else \
|
||||
echo "Installing dev requirements..." \
|
||||
&& if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
|
||||
echo "Recompiling test requirements for arm64..." \
|
||||
&& uv pip compile requirements/test/cuda.in -o requirements/test/cuda.txt --index-strategy unsafe-best-match \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
|
||||
fi \
|
||||
&& uv pip install --system -r requirements/dev.txt \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
|
||||
fi \
|
||||
|
||||
+5
-10
@@ -135,18 +135,13 @@ ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
|
||||
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
|
||||
ENV CARGO_BUILD_JOBS=4
|
||||
ENV CARGO_NET_RETRY=10
|
||||
ENV RUSTUP_MAX_RETRIES=10
|
||||
|
||||
# Build the release binary. Cargo's registry/git caches can be written by
|
||||
# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
|
||||
# keeping the cache benefit. Copy the binary out so it persists into the image
|
||||
# layer for later COPY --from=rust-build.
|
||||
RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
|
||||
--mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
|
||||
# Build the release binary. Cache cargo registry/git, and copy the binary out
|
||||
# so it persists into the image layer for later COPY --from=rust-build.
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry \
|
||||
--mount=type=cache,target=/root/.cargo/git \
|
||||
cd ${COMMON_WORKDIR}/vllm \
|
||||
&& VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \
|
||||
&& test -x /tmp/vllm-rs
|
||||
&& VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh
|
||||
|
||||
# -----------------------
|
||||
# vLLM build stages
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete
|
||||
ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.2-complete
|
||||
ARG TRITON_BRANCH="ba5c1517"
|
||||
ARG TRITON_REPO="https://github.com/ROCm/triton.git"
|
||||
ARG PYTORCH_BRANCH="8514f051" # release/2.10 as of 3/17
|
||||
@@ -104,6 +104,28 @@ ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
|
||||
ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
|
||||
ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
|
||||
|
||||
# torch profiler hotfix for 7.2.2: rebuild CLR with https://github.com/ROCm/rocm-systems/pull/5062
|
||||
# will be removed once we move to ROCm 7.2.3
|
||||
RUN apt-get update && apt-get install -y rocm-llvm-dev
|
||||
RUN pip install CppHeaderParser
|
||||
RUN git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems /tmp/rocm-systems \
|
||||
&& cd /tmp/rocm-systems \
|
||||
&& git sparse-checkout init --cone \
|
||||
&& git sparse-checkout set projects/hip projects/clr \
|
||||
&& git checkout 35e8c7bf8911862e5389509800e65fdf125412b3 \
|
||||
&& export CLR_DIR=/tmp/rocm-systems/projects/clr \
|
||||
&& export HIP_DIR=/tmp/rocm-systems/projects/hip \
|
||||
&& mkdir -p $CLR_DIR/build && cd $CLR_DIR/build \
|
||||
&& cmake \
|
||||
-DHIP_COMMON_DIR=$HIP_DIR \
|
||||
-DCMAKE_PREFIX_PATH="/opt/rocm/" \
|
||||
-DCLR_BUILD_HIP=ON \
|
||||
-DCLR_BUILD_OCL=OFF \
|
||||
-DHIP_PLATFORM=amd \
|
||||
.. \
|
||||
&& make -j$(nproc) \
|
||||
&& make install \
|
||||
&& rm -rf /tmp/rocm-systems
|
||||
|
||||
###
|
||||
### Triton Build
|
||||
@@ -237,8 +259,9 @@ ARG AITER_REPO
|
||||
ARG USE_SCCACHE
|
||||
RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
|
||||
pip install /install/*.whl
|
||||
RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO}
|
||||
RUN git clone --recursive ${AITER_REPO}
|
||||
RUN cd aiter \
|
||||
&& git checkout ${AITER_BRANCH} \
|
||||
&& git submodule update --init --recursive \
|
||||
&& pip install -r requirements.txt
|
||||
RUN pip install pyyaml && cd aiter \
|
||||
|
||||
@@ -918,41 +918,6 @@ vllm bench serve \
|
||||
|
||||
</details>
|
||||
|
||||
### Replay Timed Traces
|
||||
|
||||
<details class="admonition abstract" markdown="1">
|
||||
<summary>Show more</summary>
|
||||
|
||||
Example of how to run traces which have timing information
|
||||
with them.
|
||||
|
||||
#### Running MoonshotAI traces
|
||||
|
||||
Start the server:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3.5-2B \
|
||||
--host 127.0.0.1 --port 8000
|
||||
```
|
||||
|
||||
Run the benchmark:
|
||||
|
||||
```bash
|
||||
# Download an example trace
|
||||
# curl -L -o conversation_trace.jsonl \
|
||||
#https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl
|
||||
|
||||
vllm bench serve --model Qwen/Qwen3.5-2B \
|
||||
--dataset-name=timed_trace --num-prompts 100 --host 127.0.0.1 \
|
||||
--port 8000 --dataset-path ./conversation_trace.jsonl \
|
||||
--ignore-eos --self-timed --timed-trace-chunk-hash-size 512 \
|
||||
--timed-trace-sec-multiplier 0.001
|
||||
```
|
||||
|
||||
This will replay the first 100 lines from the trace file `conversation.jsonl`.
|
||||
|
||||
</details>
|
||||
|
||||
### 🧪 Hashing Benchmarks
|
||||
|
||||
<details class="admonition abstract" markdown="1">
|
||||
|
||||
@@ -201,45 +201,43 @@ The profiling traces generated by the continuous profiling workflow are publicly
|
||||
|
||||
The Python standard library includes
|
||||
[cProfile](https://docs.python.org/3/library/profile.html) for profiling Python
|
||||
code.
|
||||
code. vLLM includes a couple of helpers that make it easy to apply it to a section of vLLM.
|
||||
Both the `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_context` functions can be
|
||||
used to profile a section of code.
|
||||
|
||||
### Example usage - function call
|
||||
!!! note
|
||||
The `vllm.utils.profiling` helpers are deprecated and will be removed in
|
||||
`v0.21`. Please use Python's `cProfile` module directly instead.
|
||||
|
||||
If a filename is specified, the profile will be saved to that file. If no
|
||||
filename is specified, profile data can be printed to stdout.
|
||||
### Example usage - decorator
|
||||
|
||||
The first helper is a Python decorator that can be used to profile a function.
|
||||
If a filename is specified, the profile will be saved to that file. If no filename is
|
||||
specified, profile data will be printed to stdout.
|
||||
|
||||
```python
|
||||
import cProfile
|
||||
|
||||
from vllm.utils.profiling import cprofile
|
||||
|
||||
@cprofile("expensive_function.prof")
|
||||
def expensive_function():
|
||||
# some expensive code
|
||||
pass
|
||||
|
||||
|
||||
profiler = cProfile.Profile()
|
||||
profiler.runcall(expensive_function)
|
||||
profiler.dump_stats("expensive_function.prof")
|
||||
```
|
||||
|
||||
### Example usage - context manager style
|
||||
### Example Usage - context manager
|
||||
|
||||
The second helper is a context manager that can be used to profile a block of
|
||||
code. Similar to the decorator, the filename is optional.
|
||||
|
||||
```python
|
||||
import cProfile
|
||||
|
||||
from vllm.utils.profiling import cprofile_context
|
||||
|
||||
def another_function():
|
||||
# more expensive code
|
||||
pass
|
||||
|
||||
|
||||
profiler = cProfile.Profile()
|
||||
profiler.enable()
|
||||
try:
|
||||
with cprofile_context("another_function.prof"):
|
||||
another_function()
|
||||
finally:
|
||||
profiler.disable()
|
||||
profiler.dump_stats("another_function.prof")
|
||||
```
|
||||
|
||||
### Analyzing Profile Results
|
||||
|
||||
@@ -205,9 +205,8 @@ hardware and configuration.
|
||||
| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only |
|
||||
| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | DeepSeek R1 dims only |
|
||||
|
||||
> **‡** Automatic selection tries FlashAttention first. On Blackwell
|
||||
> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then
|
||||
> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.
|
||||
> **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
|
||||
> On other GPUs, FlashAttention is used as the default.
|
||||
|
||||
### Decode Backends
|
||||
|
||||
|
||||
@@ -21,6 +21,7 @@ or just on the low or high end.
|
||||
| Fusion | `PassConfig` flag | Fused operations | Default at | E2E Speedup | Fullgraph | `num_tokens` |
|
||||
| ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |
|
||||
| [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms) | `fuse_allreduce_rms` | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20% | No | Low |
|
||||
| [MiniMax QK Norm](#minimax-qk-norm-fuse_minimax_qk_norm) | `fuse_minimax_qk_norm` | Q/K variance all-reduce → Q/K RMSNorm | Off by default | 2-3% | No | Low |
|
||||
| [Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | Attention output → FP8/NVFP4 quant | Off by default | 3-7% | Yes | Always |
|
||||
| [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always |
|
||||
| [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low |
|
||||
@@ -41,6 +42,7 @@ The table below lists the quantization schemes supported by each fusion on each
|
||||
| Fusion | SM100 (Blackwell) | SM90 (Hopper) | SM89 (Ada) | SM80 (Ampere) | ROCm |
|
||||
| ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- |
|
||||
| `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — |
|
||||
| `fuse_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — |
|
||||
| `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* |
|
||||
| `fuse_attn_quant` (MLA)\* | FP8 static\*, FP8 per-group\*, NVFP4\* | FP8 static\*, FP8 per-group\* | FP8 static\*, FP8 per-group\* | — | FP8 static\* (untested) |
|
||||
| `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |
|
||||
@@ -56,6 +58,9 @@ The table below lists the quantization schemes supported by each fusion on each
|
||||
fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
|
||||
for per-backend details.
|
||||
|
||||
\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires
|
||||
tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`.
|
||||
|
||||
† `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today;
|
||||
other architectures support requires setting `PassConfig.sp_min_token_num` explicitly.
|
||||
SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`.
|
||||
@@ -186,6 +191,35 @@ If these conditions are set, the fusion is enabled automatically for optimizatio
|
||||
|
||||
- Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py)
|
||||
|
||||
### MiniMax QK Norm (`fuse_minimax_qk_norm`)
|
||||
|
||||
!!! info
|
||||
This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold:
|
||||
the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`),
|
||||
and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any
|
||||
optimization level.
|
||||
|
||||
**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the
|
||||
per-token Q/K variances before applying RMS normalization to Q and K.
|
||||
|
||||
This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion):
|
||||
`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while
|
||||
`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models.
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
vllm serve MiniMaxAI/MiniMax-M2.5 \
|
||||
--tensor-parallel-size 4 \
|
||||
--compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}'
|
||||
```
|
||||
|
||||
**Code locations.**
|
||||
|
||||
- Pass: [`vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py)
|
||||
- CUDA op: [`csrc/minimax_reduce_rms_kernel.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/minimax_reduce_rms_kernel.cu) (`minimax_allreduce_rms_qk`)
|
||||
- Workspace helper: [`vllm/model_executor/layers/mamba/lamport_workspace.py`](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/lamport_workspace.py)
|
||||
|
||||
### Sequence Parallelism (`enable_sp`)
|
||||
|
||||
**What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather,
|
||||
|
||||
@@ -31,7 +31,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct --port 8020 --kv-transfer-config '{"kv_conne
|
||||
### Proxy
|
||||
|
||||
```bash
|
||||
python examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020
|
||||
python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020
|
||||
```
|
||||
|
||||
Now you can send requests to the proxy server through port 8000.
|
||||
|
||||
@@ -15,7 +15,6 @@ vLLM currently supports the following reasoning models:
|
||||
| ------------ | ----------- | ---------------- | ----------- |
|
||||
| [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ |
|
||||
| [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ |
|
||||
| [Gemma 4 series](https://huggingface.co/google/gemma-4-26B-A4B-it) | `gemma4` | `json`, `regex` | ✅ |
|
||||
| [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ |
|
||||
| [ERNIE-4.5-VL series](https://huggingface.co/baidu/ERNIE-4.5-VL-28B-A3B-PT) | `ernie45` | `json`, `regex` | ❌ |
|
||||
| [ERNIE-4.5-21B-A3B-Thinking](https://huggingface.co/baidu/ERNIE-4.5-21B-A3B-Thinking) | `ernie45` | `json`, `regex` | ✅ |
|
||||
@@ -30,7 +29,6 @@ vLLM currently supports the following reasoning models:
|
||||
!!! note
|
||||
IBM Granite 3.2 and DeepSeek-V3.1 reasoning is disabled by default; to enable it, you must also pass `thinking=True` in your `chat_template_kwargs`.
|
||||
The reasoning feature for the Qwen3 series is enabled by default. To disable it, you must pass `enable_thinking=False` in your `chat_template_kwargs`.
|
||||
Gemma 4 reasoning is disabled by default; to enable it, pass `enable_thinking=True` in your `chat_template_kwargs` or set `reasoning_effort` (which enables it automatically).
|
||||
DeepSeek-V3.1 tool calling is supported in non-thinking mode.
|
||||
Holo2 reasoning is enabled by default. To disable it, you must also pass `thinking=False` in your `chat_template_kwargs`.
|
||||
|
||||
@@ -316,44 +314,9 @@ for output in outputs:
|
||||
print("text:", output.outputs[0].text)
|
||||
```
|
||||
|
||||
## Automatic `enable_thinking` Activation
|
||||
|
||||
Some models (such as Gemma 4, DeepSeek-V4-Pro and IBM Granite 3.2) require `enable_thinking: true` in their chat template kwargs to activate thinking mode — without it, reasoning tokens are never generated regardless of other settings.
|
||||
|
||||
When you set `reasoning_effort` in a Chat Completions request (or `reasoning.effort` in a Responses API request), vLLM automatically injects `enable_thinking` into the chat template kwargs:
|
||||
|
||||
- `reasoning_effort` = `"low"`, `"medium"`, or `"high"` → `enable_thinking = true`
|
||||
- `reasoning_effort` = `"none"` → `enable_thinking = false`
|
||||
- `reasoning_effort` not set → `enable_thinking` is not injected (preserves existing behavior)
|
||||
|
||||
This means you no longer need to manually pass `chat_template_kwargs: {"enable_thinking": true}` when using `reasoning_effort` — it is handled automatically.
|
||||
|
||||
!!! note
|
||||
If you explicitly set `enable_thinking` in `chat_template_kwargs`, your value takes priority over the automatic injection. This allows you to override the behavior if needed.
|
||||
|
||||
For models whose templates don't declare `enable_thinking` (e.g., DeepSeek R1), the injected kwarg is harmlessly filtered out by `resolve_chat_template_kwargs`.
|
||||
|
||||
### Example
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")
|
||||
|
||||
# reasoning_effort automatically enables thinking for models that need it
|
||||
response = client.chat.completions.create(
|
||||
model="google/gemma-4-26B-A4B-it",
|
||||
messages=[{"role": "user", "content": "What is 15 * 37?"}],
|
||||
reasoning_effort="high", # Automatically sets enable_thinking=true
|
||||
)
|
||||
|
||||
print(response.choices[0].message.reasoning)
|
||||
print(response.choices[0].message.content)
|
||||
```
|
||||
|
||||
## Limitations
|
||||
|
||||
- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`), Anthropic Messages API (`/v1/messages`) and the Responses API (`/v1/responses`).
|
||||
- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`).
|
||||
|
||||
## How to support a new reasoning model
|
||||
|
||||
|
||||
@@ -76,15 +76,6 @@ This guide will help you quickly get started with vLLM to perform:
|
||||
!!! note
|
||||
For more detailed instructions, including Docker, installing from source, and troubleshooting, please refer to the [vLLM on TPU documentation](https://docs.vllm.ai/projects/tpu/en/latest/).
|
||||
|
||||
=== "Ascend NPU"
|
||||
|
||||
If you are using Ascend NPUs, you can run vLLM through [vLLM Ascend](https://github.com/vllm-project/vllm-ascend), a community-maintained hardware plugin.
|
||||
|
||||
Follow the installation instructions in the [vLLM Ascend quick start](https://docs.vllm.ai/projects/ascend/en/latest/quick_start.html).
|
||||
|
||||
!!! note
|
||||
Ascend setup depends on your NPU hardware and CANN version. For supported versions, Docker images, and troubleshooting, please refer to the [vLLM Ascend documentation](https://docs.vllm.ai/projects/ascend/en/latest/).
|
||||
|
||||
=== "Apple Silicon (Mac)"
|
||||
|
||||
If you are using Apple Silicon Macs, you can use vLLM-Metal for GPU-accelerated inference via Apple's Metal framework.
|
||||
|
||||
@@ -106,8 +106,8 @@ class UrlSchemesPreprocessor(Preprocessor):
|
||||
return f"[{gh_icon} {title}]({url})"
|
||||
|
||||
markdown = "\n".join(lines)
|
||||
markdown = github_link.sub(replace_github_link, markdown)
|
||||
markdown = relative_link.sub(replace_relative_link, markdown)
|
||||
markdown = github_link.sub(replace_github_link, markdown)
|
||||
return markdown.split("\n")
|
||||
|
||||
|
||||
|
||||
@@ -299,3 +299,7 @@ Example configuration:
|
||||
### Remove softmax from PoolingParams
|
||||
|
||||
We have already removed `softmax` and `activation` from PoolingParams. Instead, use `use_activation`, since we allow `classify` and `token_classify` to use any activation function.
|
||||
|
||||
### Remove `logit_bias` and `logit_scale`
|
||||
|
||||
`logit_bias` and `logit_scale` are deprecated aliases for `logit_mean` and `logit_sigma` respectively. When using `logit_scale`, it is automatically converted to `logit_sigma = 1/logit_scale`. These deprecated parameters will be removed in v0.21.
|
||||
|
||||
@@ -428,6 +428,7 @@ th {
|
||||
| `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | |
|
||||
| `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | |
|
||||
| `JAISLMHeadModel` | Jais | `inceptionai/jais-13b`, `inceptionai/jais-13b-chat`, `inceptionai/jais-30b-v3`, `inceptionai/jais-30b-chat-v3`, etc. | | ✅︎ |
|
||||
| `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ |
|
||||
| `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ |
|
||||
| `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B-Base, Kimi-Linear-48B-A3B-Instruct | `moonshotai/Kimi-Linear-48B-A3B-Base`, `moonshotai/Kimi-Linear-48B-A3B-Instruct` | | ✅︎ |
|
||||
@@ -550,7 +551,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ |
|
||||
| `CheersForConditionalGeneration` | Cheers | T + I | `ai9stars/Cheers` | | ✅︎ |
|
||||
| `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ |
|
||||
| `Cosmos3ForConditionalGeneration` | Cosmos3 (understanding tower) | T + I<sup>E+</sup> + V<sup>E+</sup> | `nvidia/Cosmos3-Nano` | | ✅︎ |
|
||||
| `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
|
||||
| `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
|
||||
| `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -141,7 +141,6 @@ bbc5b7ede = "bbc5b7ede"
|
||||
NOOPs = "NOOPs"
|
||||
nin_shortcut = "nin_shortcut"
|
||||
cudaDevAttrMaxSharedMemoryPerBlockOptin = "cudaDevAttrMaxSharedMemoryPerBlockOptin"
|
||||
sharedMemPerBlockOptin = "sharedMemPerBlockOptin"
|
||||
|
||||
depthwise_seperable_out_channel = "depthwise_seperable_out_channel"
|
||||
pard_token = "pard_token"
|
||||
@@ -182,8 +181,6 @@ VALU = "VALU"
|
||||
# Walsh-Hadamard Transform
|
||||
wht = "wht"
|
||||
WHT = "WHT"
|
||||
# Huawei Compute Architecture for Neural Networks
|
||||
CANN = "CANN"
|
||||
|
||||
[tool.uv]
|
||||
no-build-isolation-package = ["torch"]
|
||||
|
||||
@@ -16,4 +16,3 @@ wheel
|
||||
jinja2>=3.1.6
|
||||
amdsmi==7.0.2
|
||||
timm>=1.0.17
|
||||
tilelang==0.1.10
|
||||
|
||||
@@ -21,7 +21,7 @@ nvidia-cudnn-frontend>=1.13.0,<1.19.0
|
||||
fastsafetensors >= 0.2.2
|
||||
|
||||
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
|
||||
nvidia-cutlass-dsl[cu13]==4.5.2
|
||||
nvidia-cutlass-dsl[cu13]==4.5.0
|
||||
quack-kernels>=0.3.3
|
||||
|
||||
# Tokenspeed_MLA for faster mla with spec decode
|
||||
|
||||
@@ -22,4 +22,3 @@ timm>=1.0.17
|
||||
# amd-quark: required for Quark quantization on ROCm
|
||||
# To be consistent with test_quark.py
|
||||
amd-quark>=0.8.99
|
||||
tilelang==0.1.10
|
||||
|
||||
@@ -53,12 +53,12 @@ tritonclient>=2.51.0
|
||||
grpcio==1.78.0
|
||||
grpcio-reflection==1.78.0
|
||||
|
||||
arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix decoding test
|
||||
arctic-inference == 0.1.1 # Required for suffix decoding test
|
||||
numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.2.2; platform_machine == "x86_64" # 0.2.2 contains important fixes for multi-GPU mem usage
|
||||
instanttensor>=0.1.5; platform_machine == "x86_64"
|
||||
fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
|
||||
instanttensor>=0.1.5
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
decord==0.6.0; platform_machine == "x86_64"
|
||||
# terratorch is temporarily disabled while PyPI has the `lightning` package
|
||||
|
||||
@@ -43,7 +43,6 @@ schemathesis>=3.39.15 # Required for openai schema test
|
||||
# quantization
|
||||
bitsandbytes==0.49.2
|
||||
buildkite-test-collector==0.1.9
|
||||
tilelang==0.1.10
|
||||
|
||||
genai_perf>=0.0.8
|
||||
tritonclient>=2.51.0
|
||||
|
||||
@@ -43,9 +43,7 @@ anyio==4.13.0
|
||||
# starlette
|
||||
# watchfiles
|
||||
apache-tvm-ffi==0.1.10
|
||||
# via
|
||||
# tilelang
|
||||
# xgrammar
|
||||
# via xgrammar
|
||||
arctic-inference==0.1.1
|
||||
# via -r requirements/test/rocm.in
|
||||
argcomplete==3.6.3
|
||||
@@ -131,9 +129,7 @@ click==8.3.1
|
||||
# typer
|
||||
# uvicorn
|
||||
cloudpickle==3.1.2
|
||||
# via
|
||||
# -r requirements/test/../common.txt
|
||||
# tilelang
|
||||
# via -r requirements/test/../common.txt
|
||||
colorama==0.4.6
|
||||
# via
|
||||
# perceptron
|
||||
@@ -515,8 +511,6 @@ mistral-common==1.11.2
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
# -r requirements/test/rocm.in
|
||||
ml-dtypes==0.5.4
|
||||
# via tilelang
|
||||
model-hosting-container-standards==0.1.14
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
@@ -593,7 +587,6 @@ numpy==2.2.6
|
||||
# lm-eval
|
||||
# matplotlib
|
||||
# mistral-common
|
||||
# ml-dtypes
|
||||
# mteb
|
||||
# numba
|
||||
# opencv-python-headless
|
||||
@@ -617,7 +610,6 @@ numpy==2.2.6
|
||||
# statsmodels
|
||||
# tensorizer
|
||||
# tifffile
|
||||
# tilelang
|
||||
# torchvision
|
||||
# transformers
|
||||
# tritonclient
|
||||
@@ -819,7 +811,6 @@ psutil==7.2.2
|
||||
# accelerate
|
||||
# peft
|
||||
# tensorizer
|
||||
# tilelang
|
||||
py==1.11.0
|
||||
# via pytest-forked
|
||||
py-cpuinfo==9.0.0
|
||||
@@ -1201,10 +1192,6 @@ tiktoken==0.12.0
|
||||
# gpt-oss
|
||||
# lm-eval
|
||||
# mistral-common
|
||||
tilelang==0.1.10
|
||||
# via
|
||||
# -c requirements/rocm.txt
|
||||
# -r requirements/test/rocm.in
|
||||
timm==1.0.17
|
||||
# via
|
||||
# -c requirements/rocm.txt
|
||||
@@ -1221,8 +1208,6 @@ tomli==2.4.0
|
||||
# via schemathesis
|
||||
tomli-w==1.2.0
|
||||
# via schemathesis
|
||||
torch-c-dlpack-ext==0.1.5
|
||||
# via tilelang
|
||||
tqdm==4.67.3
|
||||
# via
|
||||
# -r requirements/test/../common.txt
|
||||
@@ -1240,7 +1225,6 @@ tqdm==4.67.3
|
||||
# pqdm
|
||||
# segmentation-models-pytorch
|
||||
# sentence-transformers
|
||||
# tilelang
|
||||
# transformers
|
||||
transformers==5.5.3
|
||||
# via
|
||||
@@ -1309,7 +1293,6 @@ typing-extensions==4.15.0
|
||||
# sentence-transformers
|
||||
# sqlalchemy
|
||||
# starlette
|
||||
# tilelang
|
||||
# torch
|
||||
# typeguard
|
||||
# typing-inspection
|
||||
@@ -1376,8 +1359,6 @@ yarl==1.23.0
|
||||
# via
|
||||
# aiohttp
|
||||
# schemathesis
|
||||
z3-solver==4.15.4.0
|
||||
# via tilelang
|
||||
zipp==3.23.0
|
||||
# via importlib-metadata
|
||||
|
||||
|
||||
Generated
-42
@@ -2371,15 +2371,6 @@ version = "0.2.16"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981"
|
||||
|
||||
[[package]]
|
||||
name = "libmimalloc-sys"
|
||||
version = "0.1.49"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9"
|
||||
dependencies = [
|
||||
"cc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libredox"
|
||||
version = "0.1.14"
|
||||
@@ -2578,15 +2569,6 @@ version = "0.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b"
|
||||
|
||||
[[package]]
|
||||
name = "mimalloc"
|
||||
version = "0.1.52"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862"
|
||||
dependencies = [
|
||||
"libmimalloc-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mime"
|
||||
version = "0.3.17"
|
||||
@@ -2609,7 +2591,6 @@ version = "2.18.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "328251e58ad8e415be6198888fc207502727dc77945806421ab34f35bf012e7d"
|
||||
dependencies = [
|
||||
"indexmap 2.13.0",
|
||||
"memo-map",
|
||||
"serde",
|
||||
"serde_json",
|
||||
@@ -5627,7 +5608,6 @@ dependencies = [
|
||||
"minijinja",
|
||||
"minijinja-contrib",
|
||||
"openai-harmony",
|
||||
"paste",
|
||||
"reqwest",
|
||||
"rmp-serde",
|
||||
"serde",
|
||||
@@ -5662,7 +5642,6 @@ dependencies = [
|
||||
"educe",
|
||||
"expect-test",
|
||||
"itertools 0.14.0",
|
||||
"mimalloc",
|
||||
"native-tls",
|
||||
"serde",
|
||||
"serde_json",
|
||||
@@ -5761,25 +5740,6 @@ dependencies = [
|
||||
"prometheus-client",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-mock-engine"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"asynk-strim-attr",
|
||||
"clap",
|
||||
"futures",
|
||||
"rand 0.9.2",
|
||||
"rmpv",
|
||||
"serde",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"vllm-engine-core-client",
|
||||
"zeromq",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-reasoning-parser"
|
||||
version = "0.1.0"
|
||||
@@ -5850,7 +5810,6 @@ dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"serial_test",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
@@ -5888,7 +5847,6 @@ name = "vllm-tool-parser"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"criterion",
|
||||
"easy-ext",
|
||||
"expect-test",
|
||||
"futures",
|
||||
"openai-protocol",
|
||||
|
||||
+3
-7
@@ -6,7 +6,6 @@ members = [
|
||||
"src/llm",
|
||||
"src/managed-engine",
|
||||
"src/metrics",
|
||||
"src/mock-engine",
|
||||
"src/reasoning-parser",
|
||||
"src/server",
|
||||
"src/text",
|
||||
@@ -46,20 +45,17 @@ http-body = "1.0.1"
|
||||
itertools = "0.14.0"
|
||||
libc = "0.2.177"
|
||||
llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" }
|
||||
mimalloc = "0.1.52"
|
||||
minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
|
||||
minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls"] }
|
||||
minijinja-contrib = { version = "2.0", features = ["pycompat"] }
|
||||
native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] }
|
||||
ndarray = { version = "0.16.1", features = ["serde"] }
|
||||
openai-harmony = "0.0.8"
|
||||
openai-protocol = "1.6.0"
|
||||
parking_lot = "0.12.5"
|
||||
paste = "1.0.15"
|
||||
prometheus-client = "0.24.0"
|
||||
prometheus-client-derive-encode = "0.5.0"
|
||||
prost = "0.14.3"
|
||||
prost-types = "0.14.3"
|
||||
rand = "0.9.2"
|
||||
reasoning-parser = "1.2.2"
|
||||
reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }
|
||||
riptoken = { version = "0.3.0", default-features = false }
|
||||
@@ -69,11 +65,11 @@ rustc-hash = "1.1.0"
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
serde-json-fmt = "0.1.0"
|
||||
serde_default = "0.2.0"
|
||||
serde_json = { version = "1.0.145", features = ["arbitrary_precision", "preserve_order"] }
|
||||
serde_json = "1.0.145"
|
||||
serde_repr = "0.1.20"
|
||||
serde_tuple = "1.1.3"
|
||||
serde_with = "3.18.0"
|
||||
serial_test = { version = "3.2.0", features = ["file_locks"] }
|
||||
serial_test = "3.2.0"
|
||||
socket2 = "0.6.3"
|
||||
subenum = "1.1.3"
|
||||
task-local = "0.1.1"
|
||||
|
||||
@@ -39,9 +39,8 @@ anyhow.workspace = true
|
||||
bytes.workspace = true
|
||||
clap.workspace = true
|
||||
expect-test.workspace = true
|
||||
paste.workspace = true
|
||||
rmp-serde.workspace = true
|
||||
serial_test.workspace = true
|
||||
serial_test = { workspace = true, features = ["file_locks"] }
|
||||
tempfile.workspace = true
|
||||
tokio.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
|
||||
@@ -53,25 +53,6 @@ impl AssistantContentBlock {
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Return a copy of this block with leading and trailing whitespace trimmed from all text
|
||||
/// fields and tool call arguments, or `None` if the resulting text would be empty.
|
||||
pub fn trim(mut self) -> Option<Self> {
|
||||
match &mut self {
|
||||
Self::Text { text } | Self::Reasoning { text } => {
|
||||
let trimmed_text = text.trim();
|
||||
if trimmed_text.is_empty() {
|
||||
return None;
|
||||
} else {
|
||||
*text = trimmed_text.to_string();
|
||||
}
|
||||
}
|
||||
Self::ToolCall(call) => {
|
||||
call.arguments = call.arguments.trim().to_string();
|
||||
}
|
||||
}
|
||||
Some(self)
|
||||
}
|
||||
}
|
||||
|
||||
#[easy_ext::ext(AssistantMessageExt)]
|
||||
@@ -138,13 +119,6 @@ impl AssistantMessage {
|
||||
pub(crate) fn push_block(&mut self, block: AssistantContentBlock) {
|
||||
self.content.push(block);
|
||||
}
|
||||
|
||||
/// Return a copy of this message with leading and trailing whitespace trimmed from all text
|
||||
/// fields and tool call arguments, and with any blocks that are empty after trimming removed.
|
||||
pub fn trim(mut self) -> Self {
|
||||
self.content = self.content.into_iter().filter_map(|block| block.trim()).collect();
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Streamed chat event emitted by [`crate::ChatEventStream`].
|
||||
|
||||
@@ -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, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
+47
-119
@@ -392,13 +392,53 @@ impl MultimodalModelInfo {
|
||||
prompt_token_ids: &mut Vec<u32>,
|
||||
replacements: Vec<PromptReplacement>,
|
||||
) -> Result<Vec<PlaceholderRange>> {
|
||||
expand_prompt_token_ids(
|
||||
prompt_token_ids,
|
||||
replacements,
|
||||
self.spec.placeholder_marker_token_id,
|
||||
self.spec.placeholder_embed_token_id,
|
||||
&self.spec.placeholder_token,
|
||||
)
|
||||
let mut cursor = 0;
|
||||
let mut ranges = Vec::with_capacity(replacements.len());
|
||||
for replacement in replacements {
|
||||
if replacement.modality != Modality::Image {
|
||||
bail_multimodal!(
|
||||
"unsupported prompt replacement modality `{}`",
|
||||
replacement.modality
|
||||
);
|
||||
}
|
||||
let offset = find_next_token(
|
||||
prompt_token_ids,
|
||||
self.spec.placeholder_marker_token_id,
|
||||
cursor,
|
||||
)
|
||||
.ok_or_else(|| {
|
||||
multimodal!(
|
||||
"placeholder token `{}` was not found in tokenized prompt",
|
||||
self.spec.placeholder_token
|
||||
)
|
||||
})?;
|
||||
|
||||
if replacement.tokens.is_empty() {
|
||||
bail_multimodal!(
|
||||
"placeholder token `{}` expanded to no tokens",
|
||||
self.spec.placeholder_token
|
||||
);
|
||||
}
|
||||
let replacement_len = replacement.tokens.len();
|
||||
let replacement_tokens =
|
||||
replacement.tokens.iter().map(|&token| token as u32).collect::<Vec<_>>();
|
||||
let is_embed = {
|
||||
let mask = replacement_tokens
|
||||
.iter()
|
||||
.map(|&token| token == self.spec.placeholder_embed_token_id)
|
||||
.collect::<Vec<_>>();
|
||||
WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)?
|
||||
};
|
||||
|
||||
prompt_token_ids.splice(offset..offset + 1, replacement_tokens);
|
||||
ranges.push(PlaceholderRange {
|
||||
offset,
|
||||
length: replacement_len,
|
||||
is_embed: Some(is_embed),
|
||||
});
|
||||
cursor = offset + replacement_len;
|
||||
}
|
||||
Ok(ranges)
|
||||
}
|
||||
|
||||
/// Convert preprocessed image tensors into engine-core multimodal features.
|
||||
@@ -476,71 +516,6 @@ impl MultimodalModelInfo {
|
||||
}
|
||||
}
|
||||
|
||||
fn expand_prompt_token_ids(
|
||||
prompt_token_ids: &mut Vec<u32>,
|
||||
replacements: Vec<PromptReplacement>,
|
||||
placeholder_marker_token_id: u32,
|
||||
placeholder_embed_token_id: u32,
|
||||
placeholder_token: &str,
|
||||
) -> Result<Vec<PlaceholderRange>> {
|
||||
if replacements.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let replacement_growth = replacements.iter().fold(0usize, |total, replacement| {
|
||||
total.saturating_add(replacement.tokens.len().saturating_sub(1))
|
||||
});
|
||||
let mut expanded =
|
||||
Vec::with_capacity(prompt_token_ids.len().saturating_add(replacement_growth));
|
||||
let mut ranges = Vec::with_capacity(replacements.len());
|
||||
let mut cursor = 0usize;
|
||||
|
||||
for replacement in replacements {
|
||||
if replacement.modality != Modality::Image {
|
||||
bail_multimodal!(
|
||||
"unsupported prompt replacement modality `{}`",
|
||||
replacement.modality
|
||||
);
|
||||
}
|
||||
|
||||
let offset = find_next_token(prompt_token_ids, placeholder_marker_token_id, cursor)
|
||||
.ok_or_else(|| {
|
||||
multimodal!(
|
||||
"placeholder token `{placeholder_token}` was not found in tokenized prompt"
|
||||
)
|
||||
})?;
|
||||
|
||||
if replacement.tokens.is_empty() {
|
||||
bail_multimodal!("placeholder token `{placeholder_token}` expanded to no tokens");
|
||||
}
|
||||
|
||||
let replacement_len = replacement.tokens.len();
|
||||
let is_embed = {
|
||||
let mask = replacement
|
||||
.tokens
|
||||
.iter()
|
||||
.map(|&token| token as u32 == placeholder_embed_token_id)
|
||||
.collect::<Vec<_>>();
|
||||
WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)?
|
||||
};
|
||||
|
||||
expanded.extend_from_slice(&prompt_token_ids[cursor..offset]);
|
||||
let expanded_offset = expanded.len();
|
||||
expanded.extend(replacement.tokens.into_iter().map(|token| token as u32));
|
||||
ranges.push(PlaceholderRange {
|
||||
offset: expanded_offset,
|
||||
length: replacement_len,
|
||||
is_embed: Some(is_embed),
|
||||
});
|
||||
cursor = offset + 1;
|
||||
}
|
||||
|
||||
expanded.extend_from_slice(&prompt_token_ids[cursor..]);
|
||||
*prompt_token_ids = expanded;
|
||||
|
||||
Ok(ranges)
|
||||
}
|
||||
|
||||
/// Find `needle` in `haystack`, starting at `start`.
|
||||
///
|
||||
/// This is intentionally order-preserving rather than a global replace: each
|
||||
@@ -761,53 +736,6 @@ mod tests {
|
||||
assert!(matches!(error, Error::Multimodal(message) if message.contains("not found")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_prompt_tokens_ignores_empty_replacements() {
|
||||
let info = llama4_info();
|
||||
let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2];
|
||||
let original_prompt_token_ids = prompt_token_ids.clone();
|
||||
|
||||
let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, Vec::new()).unwrap();
|
||||
|
||||
assert!(ranges.is_empty());
|
||||
assert_eq!(prompt_token_ids, original_prompt_token_ids);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_prompt_tokens_leaves_prompt_unchanged_when_later_placeholder_missing() {
|
||||
let info = llama4_info();
|
||||
let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2];
|
||||
let original_prompt_token_ids = prompt_token_ids.clone();
|
||||
let replacements = vec![
|
||||
llama4_single_tile_replacement(),
|
||||
llama4_single_tile_replacement(),
|
||||
];
|
||||
|
||||
let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err();
|
||||
|
||||
assert!(matches!(error, Error::Multimodal(message) if message.contains("not found")));
|
||||
assert_eq!(prompt_token_ids, original_prompt_token_ids);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_prompt_tokens_errors_when_replacement_is_empty() {
|
||||
let info = llama4_info();
|
||||
let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2];
|
||||
let original_prompt_token_ids = prompt_token_ids.clone();
|
||||
let replacements = vec![PromptReplacement::sequence(
|
||||
Modality::Image,
|
||||
"<|image|>",
|
||||
Vec::new(),
|
||||
)];
|
||||
|
||||
let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err();
|
||||
|
||||
assert!(
|
||||
matches!(error, Error::Multimodal(message) if message.contains("expanded to no tokens"))
|
||||
);
|
||||
assert_eq!(prompt_token_ids, original_prompt_token_ids);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_prompt_tokens_skips_llama4_image_marker_inside_replacement() {
|
||||
let info = llama4_info();
|
||||
|
||||
@@ -15,7 +15,7 @@ use crate::Result;
|
||||
use crate::error::Error;
|
||||
use crate::event::AssistantBlockKind;
|
||||
use crate::output::generate_tool_call_id;
|
||||
use crate::parser::tool::{ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::parser::tool::{ToolCallDelta, ToolParseResult, ToolParser};
|
||||
|
||||
/// Per-stream tool parsing state.
|
||||
struct ToolState {
|
||||
@@ -57,52 +57,46 @@ impl ToolState {
|
||||
return Ok(events);
|
||||
}
|
||||
|
||||
let mut output = ToolParserOutput::default();
|
||||
let parse_result = self.parser.parse_into(&delta, &mut output);
|
||||
let parse_result = self.parser.push(&delta);
|
||||
|
||||
match parse_result {
|
||||
Ok(()) => self.process_parser_output(kind, output, &mut events)?,
|
||||
Ok(result) => self.process_parse_result(kind, result, &mut events)?,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
error = %error.as_report(),
|
||||
"tool parser failed; falling back to plain text deltas"
|
||||
);
|
||||
// Permanently mark this parser as failed.
|
||||
// TODO: we may consider recovering from parsing errors in the future.
|
||||
self.parser_failed = true;
|
||||
|
||||
// On parsing failure, we still apply the partial parser output if any, but we close
|
||||
// any open tool calls and emit the remaining buffered text as a plain-text delta to
|
||||
// preserve as much of the output as possible.
|
||||
self.process_parser_output(kind, output, &mut events)?;
|
||||
if !self.parser_failed {
|
||||
warn!(
|
||||
error = %error.as_report(),
|
||||
"tool parser failed; falling back to plain text deltas"
|
||||
);
|
||||
self.parser_failed = true;
|
||||
}
|
||||
self.open_call_index = None;
|
||||
push_text_delta(&mut events, kind, self.parser.reset());
|
||||
events.push(AssistantEvent::TextDelta { kind, delta });
|
||||
}
|
||||
}
|
||||
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
/// Apply one parsed tool output to the current stream state.
|
||||
fn process_parser_output(
|
||||
/// Apply one parsed tool result to the current stream state.
|
||||
fn process_parse_result(
|
||||
&mut self,
|
||||
kind: AssistantBlockKind,
|
||||
output: ToolParserOutput,
|
||||
result: ToolParseResult,
|
||||
events: &mut Vec<AssistantEvent>,
|
||||
) -> Result<()> {
|
||||
// When we are not currently streaming a tool call, preserve plain
|
||||
// text first and then surface any new tool call items.
|
||||
if self.open_call_index.is_none() {
|
||||
push_text_delta(events, kind, output.normal_text);
|
||||
self.process_tool_items(output.calls, events)?;
|
||||
push_text_delta(events, kind, result.normal_text);
|
||||
self.process_tool_items(result.calls, events)?;
|
||||
} else {
|
||||
// Once a tool call is open, prioritize tool deltas first. If the
|
||||
// parser emits normal text again, close the tool call and resume
|
||||
// plain text output.
|
||||
self.process_tool_items(output.calls, events)?;
|
||||
if !output.normal_text.is_empty() {
|
||||
self.process_tool_items(result.calls, events)?;
|
||||
if !result.normal_text.is_empty() {
|
||||
self.open_call_index = None;
|
||||
push_text_delta(events, kind, output.normal_text);
|
||||
push_text_delta(events, kind, result.normal_text);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
@@ -164,8 +158,8 @@ impl ToolState {
|
||||
}
|
||||
|
||||
match self.parser.finish() {
|
||||
Ok(output) => {
|
||||
self.process_parser_output(AssistantBlockKind::Text, output, &mut events)?
|
||||
Ok(result) => {
|
||||
self.process_parse_result(AssistantBlockKind::Text, result, &mut events)?
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(
|
||||
@@ -271,24 +265,17 @@ mod tests {
|
||||
use crate::error::Error;
|
||||
use crate::event::{AssistantBlockKind, AssistantMessageExt as _};
|
||||
use crate::output::structured::structured_chat_event_stream;
|
||||
use crate::parser::tool::{
|
||||
DeepSeekV4ToolParser, ToolParser, ToolParserError, ToolParserOutput,
|
||||
};
|
||||
use crate::parser::tool::{ToolParseResult, ToolParser, ToolParserError};
|
||||
use crate::request::ChatTool;
|
||||
use crate::stream::{ChatEventStream, CollectedAssistantMessage};
|
||||
use crate::stream::ChatEventStream;
|
||||
|
||||
struct FailingParser {
|
||||
fail_next: bool,
|
||||
buffered: String,
|
||||
}
|
||||
|
||||
struct ScriptedParser {
|
||||
push_outputs: Vec<ToolParserOutput>,
|
||||
finish_output: ToolParserOutput,
|
||||
}
|
||||
|
||||
struct PartialThenFailParser {
|
||||
buffered: String,
|
||||
push_results: Vec<ToolParseResult>,
|
||||
finish_result: ToolParseResult,
|
||||
}
|
||||
|
||||
impl ToolParser for FailingParser {
|
||||
@@ -296,14 +283,10 @@ mod tests {
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self {
|
||||
fail_next: false,
|
||||
buffered: String::new(),
|
||||
}))
|
||||
Ok(Box::new(Self { fail_next: false }))
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, _output: &mut ToolParserOutput) -> Result<()> {
|
||||
self.buffered.push_str(chunk);
|
||||
fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
|
||||
if self.fail_next {
|
||||
self.fail_next = false;
|
||||
return Err(ToolParserError::ParsingFailed {
|
||||
@@ -311,16 +294,7 @@ mod tests {
|
||||
});
|
||||
}
|
||||
|
||||
self.buffered.clear();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
Ok(ToolParserOutput::default())
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
std::mem::take(&mut self.buffered)
|
||||
Ok(ToolParseResult::default())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -330,215 +304,18 @@ mod tests {
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self {
|
||||
push_outputs: Vec::new(),
|
||||
finish_output: ToolParserOutput::default(),
|
||||
push_results: Vec::new(),
|
||||
finish_result: ToolParseResult::default(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
let mut next = self.push_outputs.pop().unwrap_or_default();
|
||||
output.normal_text.push_str(&next.normal_text);
|
||||
output.calls.append(&mut next.calls);
|
||||
Ok(())
|
||||
fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
|
||||
Ok(self.push_results.pop().unwrap_or_default())
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
Ok(std::mem::take(&mut self.finish_output))
|
||||
fn finish(&mut self) -> Result<ToolParseResult> {
|
||||
Ok(std::mem::take(&mut self.finish_result))
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
String::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolParser for PartialThenFailParser {
|
||||
fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self {
|
||||
buffered: String::new(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
output.calls.extend([
|
||||
crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some("get_weather".to_string()),
|
||||
arguments: String::new(),
|
||||
},
|
||||
crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: None,
|
||||
arguments: r#"{"location":"SF"}"#.to_string(),
|
||||
},
|
||||
]);
|
||||
self.buffered.push_str(" trailing text");
|
||||
Err(ToolParserError::ParsingFailed {
|
||||
message: "boom".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
Ok(ToolParserOutput::default())
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
std::mem::take(&mut self.buffered)
|
||||
}
|
||||
}
|
||||
|
||||
fn deepseek_v4_test_tools() -> Vec<ChatTool> {
|
||||
vec![
|
||||
ChatTool {
|
||||
name: "get_weather".to_string(),
|
||||
description: None,
|
||||
parameters: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": { "type": "string" }
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
ChatTool {
|
||||
name: "add".to_string(),
|
||||
description: None,
|
||||
parameters: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"x": { "type": "integer" },
|
||||
"y": { "type": "integer" }
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
async fn collect_deepseek_v4_message(chunks: Vec<String>) -> CollectedAssistantMessage {
|
||||
let events = chunks
|
||||
.into_iter()
|
||||
.map(|delta| {
|
||||
Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta,
|
||||
})
|
||||
})
|
||||
.chain(std::iter::once(Ok(ContentEvent::Done {
|
||||
prompt_token_count: 1,
|
||||
output_token_count: 1,
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
})));
|
||||
let parser = DeepSeekV4ToolParser::create(&deepseek_v4_test_tools()).unwrap();
|
||||
let assistant_events = tool_event_stream(stream::iter(events), Some(parser));
|
||||
let chat_events = structured_chat_event_stream(assistant_events);
|
||||
|
||||
ChatEventStream::new("req_deepseek_v4".to_string(), Box::pin(chat_events))
|
||||
.collect_message()
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn message_tool_projection(
|
||||
message: &CollectedAssistantMessage,
|
||||
) -> (String, Vec<(String, serde_json::Value)>) {
|
||||
(
|
||||
message.message.text(),
|
||||
message
|
||||
.message
|
||||
.tool_calls()
|
||||
.map(|call| {
|
||||
(
|
||||
call.name.clone(),
|
||||
serde_json::from_str(&call.arguments).unwrap(),
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_parser_error_preserves_partial_output_and_flushes_buffer() {
|
||||
let events = stream::iter(vec![
|
||||
Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "ignored".to_string(),
|
||||
}),
|
||||
Ok(ContentEvent::Done {
|
||||
prompt_token_count: 1,
|
||||
output_token_count: 1,
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let events = tool_event_stream(
|
||||
events,
|
||||
Some(Box::new(PartialThenFailParser {
|
||||
buffered: String::new(),
|
||||
})),
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
&events[0],
|
||||
AssistantEvent::ToolCallStart { name, .. } if name == "get_weather"
|
||||
));
|
||||
assert!(matches!(
|
||||
&events[1],
|
||||
AssistantEvent::ToolCallArgumentsDelta { delta } if delta == r#"{"location":"SF"}"#
|
||||
));
|
||||
assert_eq!(
|
||||
events[2],
|
||||
AssistantEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: " trailing text".to_string(),
|
||||
}
|
||||
);
|
||||
assert!(matches!(events[3], AssistantEvent::Done { .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn real_buffered_parser_error_matches_streaming_and_non_streaming() {
|
||||
let prefix = "I will check both.\n";
|
||||
let first_tool_call = concat!(
|
||||
"<｜DSML｜tool_calls>\n",
|
||||
"<｜DSML｜invoke name=\"get_weather\">\n",
|
||||
"<｜DSML｜parameter name=\"location\" string=\"true\">Tokyo</｜DSML｜parameter>\n",
|
||||
"</｜DSML｜invoke>",
|
||||
);
|
||||
let malformed_second_tool_call = concat!(
|
||||
"\n<｜DSML｜invoke name=\"add\">\n",
|
||||
"not a parameter\n",
|
||||
"</｜DSML｜invoke>\n",
|
||||
"</｜DSML｜tool_calls>",
|
||||
);
|
||||
let streaming_chunks = vec![
|
||||
prefix.to_string(),
|
||||
first_tool_call.to_string(),
|
||||
malformed_second_tool_call.to_string(),
|
||||
];
|
||||
let full_output = streaming_chunks.concat();
|
||||
|
||||
let streaming = collect_deepseek_v4_message(streaming_chunks).await;
|
||||
let non_streaming = collect_deepseek_v4_message(vec![full_output]).await;
|
||||
|
||||
let expected = (
|
||||
format!("{prefix}{malformed_second_tool_call}"),
|
||||
vec![(
|
||||
"get_weather".to_string(),
|
||||
serde_json::json!({ "location": "Tokyo" }),
|
||||
)],
|
||||
);
|
||||
assert_eq!(message_tool_projection(&streaming), expected);
|
||||
assert_eq!(message_tool_projection(&non_streaming), expected);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -564,15 +341,10 @@ mod tests {
|
||||
}),
|
||||
]);
|
||||
|
||||
let collected = tool_event_stream(
|
||||
events,
|
||||
Some(Box::new(FailingParser {
|
||||
fail_next: true,
|
||||
buffered: String::new(),
|
||||
})),
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
let collected =
|
||||
tool_event_stream(events, Some(Box::new(FailingParser { fail_next: true })))
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
|
||||
let events = collected
|
||||
.into_iter()
|
||||
@@ -643,18 +415,12 @@ mod tests {
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
let events = tool_event_stream(
|
||||
events,
|
||||
Some(Box::new(FailingParser {
|
||||
fail_next: false,
|
||||
buffered: String::new(),
|
||||
})),
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap();
|
||||
let events = tool_event_stream(events, Some(Box::new(FailingParser { fail_next: false })))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
@@ -702,7 +468,7 @@ mod tests {
|
||||
]);
|
||||
|
||||
let parser = ScriptedParser {
|
||||
push_outputs: vec![ToolParserOutput {
|
||||
push_results: vec![ToolParseResult {
|
||||
normal_text: String::new(),
|
||||
calls: vec![
|
||||
crate::parser::tool::ToolCallDelta {
|
||||
@@ -717,14 +483,14 @@ mod tests {
|
||||
},
|
||||
],
|
||||
}],
|
||||
finish_output: ToolParserOutput::default(),
|
||||
finish_result: ToolParseResult::default(),
|
||||
};
|
||||
|
||||
let err = tool_event_stream(events, Some(Box::new(parser)))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.find_map(|output| output.err())
|
||||
.find_map(|result| result.err())
|
||||
.expect("expected invariant error");
|
||||
|
||||
assert!(matches!(err, Error::ToolCallStreamInvariant { .. }));
|
||||
@@ -748,8 +514,8 @@ mod tests {
|
||||
]);
|
||||
|
||||
let parser = ScriptedParser {
|
||||
push_outputs: vec![
|
||||
ToolParserOutput {
|
||||
push_results: vec![
|
||||
ToolParseResult {
|
||||
normal_text: String::new(),
|
||||
calls: vec![crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
@@ -757,11 +523,11 @@ mod tests {
|
||||
arguments: "}".to_string(),
|
||||
}],
|
||||
},
|
||||
ToolParserOutput {
|
||||
ToolParseResult {
|
||||
normal_text: "plain text".to_string(),
|
||||
calls: Vec::new(),
|
||||
},
|
||||
ToolParserOutput {
|
||||
ToolParseResult {
|
||||
normal_text: String::new(),
|
||||
calls: vec![crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
@@ -770,14 +536,14 @@ mod tests {
|
||||
}],
|
||||
},
|
||||
],
|
||||
finish_output: ToolParserOutput::default(),
|
||||
finish_result: ToolParseResult::default(),
|
||||
};
|
||||
|
||||
let err = tool_event_stream(events, Some(Box::new(parser)))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.find_map(|output| output.err())
|
||||
.find_map(|result| result.err())
|
||||
.expect("expected invariant error");
|
||||
|
||||
assert!(matches!(
|
||||
@@ -807,7 +573,7 @@ mod tests {
|
||||
]);
|
||||
|
||||
let parser = ScriptedParser {
|
||||
push_outputs: vec![ToolParserOutput {
|
||||
push_results: vec![ToolParseResult {
|
||||
normal_text: String::new(),
|
||||
calls: vec![
|
||||
crate::parser::tool::ToolCallDelta {
|
||||
@@ -822,7 +588,7 @@ mod tests {
|
||||
},
|
||||
],
|
||||
}],
|
||||
finish_output: ToolParserOutput::default(),
|
||||
finish_result: ToolParseResult::default(),
|
||||
};
|
||||
|
||||
let events = tool_event_stream(events, Some(Box::new(parser)))
|
||||
|
||||
@@ -4,10 +4,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,
|
||||
Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, KimiK2ToolParser,
|
||||
Llama3JsonToolParser, MinimaxM2ToolParser, MistralToolParser, Qwen3CoderToolParser,
|
||||
Qwen3XmlToolParser, ToolCallDelta, ToolParseResult, ToolParser, ToolParserError,
|
||||
};
|
||||
|
||||
use crate::parser::ParserFactory;
|
||||
@@ -23,7 +22,6 @@ pub mod names {
|
||||
pub const GLM47: &str = "glm47";
|
||||
pub const GEMMA4: &str = "gemma4";
|
||||
pub const HERMES: &str = "hermes";
|
||||
pub const HY_V3: &str = "hy_v3";
|
||||
pub const KIMI_K2: &str = "kimi_k2";
|
||||
pub const LLAMA3_JSON: &str = "llama3_json";
|
||||
pub const LLAMA4_JSON: &str = "llama4_json";
|
||||
@@ -61,7 +59,6 @@ impl ToolParserFactory {
|
||||
.register_parser::<Glm47MoeToolParser>(names::GLM47)
|
||||
.register_parser::<Gemma4ToolParser>(names::GEMMA4)
|
||||
.register_parser::<HermesToolParser>(names::HERMES)
|
||||
.register_parser::<HyV3ToolParser>(names::HY_V3)
|
||||
.register_parser::<KimiK2ToolParser>(names::KIMI_K2)
|
||||
.register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
|
||||
.register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
|
||||
@@ -78,8 +75,6 @@ impl ToolParserFactory {
|
||||
.register_pattern("qwen3.5", names::QWEN3_CODER)
|
||||
.register_pattern("qwen", names::QWEN3_XML)
|
||||
.register_pattern("hermes", names::HERMES)
|
||||
.register_pattern("hy3", names::HY_V3)
|
||||
.register_pattern("hy_v3", names::HY_V3)
|
||||
.register_pattern("llama-4", names::LLAMA4_JSON)
|
||||
.register_pattern("llama-3.2", names::LLAMA3_JSON)
|
||||
.register_pattern("llama-3.1", names::LLAMA3_JSON)
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use vllm_tool_parser::Result;
|
||||
|
||||
use super::{ToolParser, ToolParserFactory, ToolParserOutput, names};
|
||||
use super::{ToolParseResult, ToolParser, ToolParserFactory, names};
|
||||
use crate::Error;
|
||||
use crate::request::ChatTool;
|
||||
|
||||
@@ -18,16 +18,8 @@ impl ToolParser for FakeToolParser {
|
||||
true
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, _chunk: &str, _output: &mut ToolParserOutput) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
Ok(ToolParserOutput::default())
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
String::new()
|
||||
fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
|
||||
Ok(ToolParseResult::default())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -149,10 +141,6 @@ fn factory_new_resolves_default_patterns() {
|
||||
factory.resolve_name_for_model("NousResearch/Hermes-3-Llama-3.1-8B"),
|
||||
Some(names::HERMES)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("tencent/Hy3-preview"),
|
||||
Some(names::HY_V3)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("MiniMax/MiniMax-M2-01"),
|
||||
Some(names::MINIMAX_M2)
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use minijinja::value::{Kwargs, ViaDeserialize};
|
||||
use minijinja::{Error as MinijinjaError, ErrorKind, Value};
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value as JsonValue;
|
||||
use serde_json::{self, Value as JsonValue};
|
||||
use serde_json_fmt::{JsonFormat, JsonSyntaxError};
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
@@ -13,7 +13,7 @@ use thiserror_ext::AsReport;
|
||||
/// - extra kwargs such as `ensure_ascii`, `separators`, and `sort_keys`
|
||||
/// - Python-style `indent` handling
|
||||
pub(super) fn hf_tojson_filter(
|
||||
ViaDeserialize(value): ViaDeserialize<JsonValue>,
|
||||
value: Value,
|
||||
kwargs: Kwargs,
|
||||
) -> std::result::Result<Value, MinijinjaError> {
|
||||
let ensure_ascii = kwargs.get::<Option<bool>>("ensure_ascii")?.unwrap_or(false);
|
||||
@@ -30,11 +30,18 @@ pub(super) fn hf_tojson_filter(
|
||||
|
||||
kwargs.assert_all_used()?;
|
||||
|
||||
let json_value: serde_json::Value = serde_json::to_value(&value).map_err(|e| {
|
||||
MinijinjaError::new(
|
||||
ErrorKind::InvalidOperation,
|
||||
format!("Failed to convert to JSON value: {e}"),
|
||||
)
|
||||
})?;
|
||||
|
||||
let json_str = {
|
||||
let value_to_serialize = if sort_keys {
|
||||
&sort_json_keys(&value)
|
||||
&sort_json_keys(&json_value)
|
||||
} else {
|
||||
&value
|
||||
&json_value
|
||||
};
|
||||
|
||||
build_json_format(indent, separators.0, separators.1, ensure_ascii)?
|
||||
@@ -207,14 +214,6 @@ mod tests {
|
||||
assert_eq!(rendered, "{\"x\":[1,2]}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tojson_preserves_arbitrary_precision_number_spelling() {
|
||||
let payload = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
|
||||
let rendered = render("{{ payload|tojson }}", payload);
|
||||
|
||||
assert_eq!(rendered, "{\"x\": 2, \"y\": 1.00}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tojson_supports_negative_indent_as_newline_only() {
|
||||
let rendered = render("{{ payload|tojson(indent=-1) }}", json!([1, 2]));
|
||||
|
||||
@@ -1,541 +0,0 @@
|
||||
//! Text-level roundtrip tests for the real chat-template and output-processor pairing.
|
||||
//!
|
||||
//! The invariant under test is that a structured assistant message rendered as history can be
|
||||
//! parsed from the generated assistant completion and then rendered back to the exact same
|
||||
//! assistant-completion text.
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context as _, Result, bail, ensure};
|
||||
use futures::{Stream, StreamExt as _, stream};
|
||||
use serde_json_fmt::JsonFormat as JsonFmt;
|
||||
use serial_test::file_serial;
|
||||
use vllm_chat::{
|
||||
AssistantContentBlock, AssistantMessage, AssistantMessageExt as _, AssistantToolCall,
|
||||
ChatEvent, ChatMessage, ChatRequest, ChatRole, ChatTool, ChatToolChoice, FinishReason,
|
||||
GenerationPromptMode, LoadModelBackendsOptions, NewChatOutputProcessorOptions, ParserSelection,
|
||||
RendererSelection, load_model_backends,
|
||||
};
|
||||
use vllm_text::{DecodedTextEvent, Finished, Prompt};
|
||||
|
||||
/// One model/parser configuration used to run the fixed roundtrip fixtures.
|
||||
struct RoundtripCase {
|
||||
/// Hugging Face model id resolved through the production backend loader.
|
||||
model_id: &'static str,
|
||||
/// Final assistant-history suffix rendered by the chat template but not
|
||||
/// generated by the model body consumed by the output processor.
|
||||
// TODO: we should adopt `ContinueFinalAssistant` mode to naturally handle this.
|
||||
assistant_stop_suffix: &'static str,
|
||||
/// Tool parser selection used by the output processor.
|
||||
tool_call_parser: ParserSelection,
|
||||
/// Reasoning parser selection used by the output processor.
|
||||
reasoning_parser: ParserSelection,
|
||||
/// JSON formatting expected after this model's template has materialized
|
||||
/// tool-call arguments.
|
||||
json_fmt: JsonFmt,
|
||||
}
|
||||
|
||||
impl RoundtripCase {
|
||||
/// Qwen3 XML tool-call format with `qwen3` reasoning tags.
|
||||
fn qwen3() -> Self {
|
||||
Self {
|
||||
model_id: "Qwen/Qwen3-0.6B",
|
||||
assistant_stop_suffix: "<|im_end|>\n",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
json_fmt: spaced_json_fmt(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Qwen3.5 coder-style JSON tool-call format with `qwen3` reasoning tags.
|
||||
fn qwen35() -> Self {
|
||||
Self {
|
||||
model_id: "Qwen/Qwen3.5-4B",
|
||||
assistant_stop_suffix: "<|im_end|>\n",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
|
||||
/// MiniMax M2.5 XML invoke format with `<think>` reasoning tags.
|
||||
fn minimax_m25() -> Self {
|
||||
Self {
|
||||
model_id: "MiniMaxAI/MiniMax-M2.5",
|
||||
assistant_stop_suffix: "[e~[\n",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
|
||||
/// DeepSeek V4 DSML tool-call format.
|
||||
fn deepseek_v4() -> Self {
|
||||
Self {
|
||||
model_id: "deepseek-ai/DeepSeek-V4-Flash",
|
||||
assistant_stop_suffix: "<｜end▁of▁sentence｜>",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
|
||||
/// GLM-4.7 XML-like argument format with `<think>` reasoning tags.
|
||||
fn glm47() -> Self {
|
||||
Self {
|
||||
model_id: "zai-org/GLM-4.7-Flash",
|
||||
assistant_stop_suffix: "",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Kimi K2.5 tool-call format with `<think>` reasoning tags.
|
||||
#[allow(dead_code)]
|
||||
fn kimi_k25() -> Self {
|
||||
Self {
|
||||
model_id: "moonshotai/Kimi-K2.5",
|
||||
assistant_stop_suffix: "<|im_end|>",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
json_fmt: spaced_json_fmt(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
macro_rules! roundtrip_tests {
|
||||
($($case:ident => [$($fixture:ident),* $(,)?]),+ $(,)?) => {
|
||||
paste::paste! {
|
||||
$(
|
||||
$(
|
||||
#[tokio::test]
|
||||
#[file_serial([<hf_ $case>])]
|
||||
async fn [<roundtrip_ $case _ $fixture>]() -> Result<()> {
|
||||
[<run_roundtrip_ $fixture>](RoundtripCase::$case()).await
|
||||
}
|
||||
)*
|
||||
)+
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
roundtrip_tests! {
|
||||
qwen3 => [reasoning_and_content, tool_call_mix],
|
||||
qwen35 => [reasoning_and_content, tool_call_mix],
|
||||
minimax_m25 => [reasoning_and_content, tool_call_mix],
|
||||
deepseek_v4 => [reasoning_and_content, tool_call_mix],
|
||||
glm47 => [reasoning_and_content, tool_call_mix],
|
||||
|
||||
// Note: Kimi K2.5 strips the reasoning content in history.
|
||||
// TODO: we don't respect model-generated tool call id now so `tool_call_mix` cannot pass.
|
||||
// kimi_k25 => [tool_call_mix],
|
||||
}
|
||||
|
||||
/// Run the fixed reasoning+content fixture for one model/parser case.
|
||||
async fn run_roundtrip_reasoning_and_content(case: RoundtripCase) -> 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(),
|
||||
);
|
||||
let expected_reasoning = "Need compute 2 + 2 directly.";
|
||||
let expected_text = "The answer is 4.";
|
||||
|
||||
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?;
|
||||
|
||||
assert_eq!(
|
||||
result.parsed_message.reasoning().as_deref().map(str::trim),
|
||||
Some(expected_reasoning)
|
||||
);
|
||||
assert_eq!(result.parsed_message.text().trim(), expected_text);
|
||||
assert_eq!(result.parsed_message.tool_calls().count(), 0);
|
||||
|
||||
assert_eq!(
|
||||
result.rerendered_closed_completion,
|
||||
result.closed_completion
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run the fixed reasoning+multiple-tools fixture for one model/parser case.
|
||||
async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
|
||||
let backends = load_roundtrip_backends(&case).await?;
|
||||
let request = roundtrip_request(
|
||||
"roundtrip-reasoning-tools",
|
||||
vec![ChatMessage::text(
|
||||
ChatRole::User,
|
||||
"Check Shanghai weather and add 1.00 plus 2.",
|
||||
)],
|
||||
test_tools(),
|
||||
);
|
||||
let expected_reasoning = "Need call the weather and add tools.";
|
||||
let expected_text = "I will call the tools.";
|
||||
|
||||
let result = run_roundtrip(
|
||||
&case,
|
||||
&backends,
|
||||
&request,
|
||||
AssistantMessage {
|
||||
content: vec![
|
||||
AssistantContentBlock::Reasoning {
|
||||
text: expected_reasoning.to_string(),
|
||||
},
|
||||
AssistantContentBlock::Text {
|
||||
text: expected_text.to_string(),
|
||||
},
|
||||
AssistantContentBlock::ToolCall(AssistantToolCall {
|
||||
id: "functions.get_weather:0".to_string(),
|
||||
name: "get_weather".to_string(),
|
||||
arguments: r#"{"location":"Shanghai"}"#.to_string(),
|
||||
}),
|
||||
AssistantContentBlock::ToolCall(AssistantToolCall {
|
||||
id: "functions.add:1".to_string(),
|
||||
name: "add".to_string(),
|
||||
// Intentionally use a non-lexical order of keys and a different number
|
||||
// formatting style to verify text-level fidelity of the roundtrip.
|
||||
arguments: r#"{"y":1.00,"x":2}"#.to_string(),
|
||||
}),
|
||||
],
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
assert_eq!(
|
||||
result.parsed_message.reasoning().as_deref().map(str::trim),
|
||||
Some(expected_reasoning)
|
||||
);
|
||||
assert_eq!(result.parsed_message.text().trim(), expected_text);
|
||||
|
||||
let tool_calls = result.parsed_message.tool_calls().collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
tool_calls.len(),
|
||||
2,
|
||||
"parsed message: {:#?}",
|
||||
result.parsed_message
|
||||
);
|
||||
assert_eq!(tool_calls[0].name, "get_weather");
|
||||
assert_eq!(
|
||||
tool_calls[0].arguments,
|
||||
expected_arguments(&case, r#"{"location": "Shanghai"}"#)?,
|
||||
);
|
||||
assert_eq!(tool_calls[1].name, "add");
|
||||
assert_eq!(
|
||||
tool_calls[1].arguments,
|
||||
expected_arguments(&case, r#"{"y": 1.00, "x": 2}"#)?,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
result.rerendered_closed_completion,
|
||||
result.closed_completion
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Compact JSON argument formatting used by JSON-native parsers/renderers.
|
||||
fn compact_json_fmt() -> JsonFmt {
|
||||
JsonFmt::new()
|
||||
}
|
||||
|
||||
/// Python `json.dumps`-style compact formatting with a space after commas and
|
||||
/// colons.
|
||||
fn spaced_json_fmt() -> JsonFmt {
|
||||
JsonFmt::new()
|
||||
.comma(", ")
|
||||
.expect("literal comma separator is valid JSON")
|
||||
.colon(": ")
|
||||
.expect("literal colon separator is valid JSON")
|
||||
}
|
||||
|
||||
/// Parse and format expected tool-call arguments from raw JSON text.
|
||||
/// Pass in a raw JSON string instead of a structured value to ensure the exact precision and
|
||||
/// formatting of numbers are preserved.
|
||||
fn expected_arguments(case: &RoundtripCase, raw_json: &str) -> Result<String> {
|
||||
let value: serde_json::Value =
|
||||
serde_json::from_str(raw_json).context("invalid expected tool-call arguments")?;
|
||||
|
||||
case.json_fmt
|
||||
.format_to_string(&value)
|
||||
.context("failed to format expected tool-call arguments")
|
||||
}
|
||||
|
||||
/// Load the real model chat/text backend for one roundtrip case.
|
||||
async fn load_roundtrip_backends(case: &RoundtripCase) -> Result<vllm_chat::LoadedModelBackends> {
|
||||
load_model_backends(
|
||||
case.model_id,
|
||||
LoadModelBackendsOptions {
|
||||
renderer: RendererSelection::Auto,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.with_context(|| format!("failed to load HF model files for {}", case.model_id))
|
||||
}
|
||||
|
||||
/// Roundtrip artifacts needed for semantic and exact-text assertions.
|
||||
struct RoundtripResult {
|
||||
/// Final assistant message reconstructed by the output processor.
|
||||
parsed_message: AssistantMessage,
|
||||
/// Assistant-completion suffix cut from rendering the expected assistant as
|
||||
/// history.
|
||||
closed_completion: String,
|
||||
/// Assistant-completion suffix cut after rendering the parsed assistant
|
||||
/// back as history.
|
||||
rerendered_closed_completion: String,
|
||||
}
|
||||
|
||||
/// Render, parse, and rerender one assistant turn through the production
|
||||
/// renderer/output-processor boundary.
|
||||
async fn run_roundtrip(
|
||||
case: &RoundtripCase,
|
||||
backends: &vllm_chat::LoadedModelBackends,
|
||||
request: &ChatRequest,
|
||||
assistant: AssistantMessage,
|
||||
) -> Result<RoundtripResult> {
|
||||
let renderer = backends.chat_backend.chat_renderer();
|
||||
let (prompt, closed_completion_text) =
|
||||
render_closed_completion(renderer.as_ref(), request, &assistant)?;
|
||||
let completion_body = closed_completion_text
|
||||
.strip_suffix(case.assistant_stop_suffix)
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"closed assistant completion did not end with {:?}: {:?}",
|
||||
case.assistant_stop_suffix, closed_completion_text
|
||||
)
|
||||
})?;
|
||||
|
||||
let parsed_message =
|
||||
parse_completion(case, backends, request, &prompt, completion_body).await?;
|
||||
let (_, rerendered_closed_completion) =
|
||||
render_closed_completion(renderer.as_ref(), request, &parsed_message)?;
|
||||
|
||||
Ok(RoundtripResult {
|
||||
parsed_message,
|
||||
closed_completion: closed_completion_text,
|
||||
rerendered_closed_completion,
|
||||
})
|
||||
}
|
||||
|
||||
/// Render `history` as a production prompt and `history + assistant` as closed
|
||||
/// history, then return the production prompt and assistant-completion suffix.
|
||||
fn render_closed_completion(
|
||||
renderer: &dyn vllm_chat::ChatRenderer,
|
||||
base_request: &ChatRequest,
|
||||
assistant: &AssistantMessage,
|
||||
) -> Result<(String, String)> {
|
||||
let mut prompt_request = base_request.clone();
|
||||
prompt_request.chat_options.generation_prompt_mode = GenerationPromptMode::StartNewAssistant;
|
||||
let prompt = render_text(renderer, &prompt_request).context("failed to render prompt")?;
|
||||
|
||||
let mut full_request = base_request.clone();
|
||||
full_request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
|
||||
full_request.messages.push(ChatMessage::from(assistant.clone()));
|
||||
let full = render_text(renderer, &full_request).context("failed to render full prompt")?;
|
||||
|
||||
ensure!(
|
||||
full.starts_with(&prompt),
|
||||
"full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
|
||||
);
|
||||
let completion = full[prompt.len()..].to_string();
|
||||
|
||||
Ok((prompt, completion))
|
||||
}
|
||||
|
||||
/// Render one chat request and require a text prompt.
|
||||
fn render_text(renderer: &dyn vllm_chat::ChatRenderer, request: &ChatRequest) -> Result<String> {
|
||||
match renderer.render(request)?.prompt {
|
||||
Prompt::Text(text) => Ok(text),
|
||||
other => bail!("roundtrip tests expect text prompts, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Feed one rendered assistant completion body into the real output processor
|
||||
/// and collect its terminal assistant message.
|
||||
async fn parse_completion(
|
||||
case: &RoundtripCase,
|
||||
backends: &vllm_chat::LoadedModelBackends,
|
||||
base_request: &ChatRequest,
|
||||
prompt: &str,
|
||||
completion_body: &str,
|
||||
) -> Result<AssistantMessage> {
|
||||
let tokenizer = backends.text_backend.tokenizer();
|
||||
let prompt_token_ids = tokenizer
|
||||
.encode(prompt, base_request.add_special_tokens)
|
||||
.context("failed to encode rendered prompt")?;
|
||||
|
||||
let mut request = base_request.clone();
|
||||
let processor = backends.chat_backend.new_chat_output_processor(
|
||||
&mut request,
|
||||
NewChatOutputProcessorOptions {
|
||||
tool_call_parser: &case.tool_call_parser,
|
||||
reasoning_parser: &case.reasoning_parser,
|
||||
},
|
||||
)?;
|
||||
|
||||
let decoded = decoded_completion_stream(prompt_token_ids, completion_body);
|
||||
let mut events = processor.process(decoded)?;
|
||||
|
||||
while let Some(event) = events.next().await {
|
||||
if let ChatEvent::Done { message, .. } = event? {
|
||||
// TODO: currently our parsers are not very strict about preserving or trimming
|
||||
// whitespace, so we trim here to avoid roundtrip failures due to
|
||||
// insignificant whitespace differences. However, this may hurt token-level
|
||||
// fidelity so we should consider improving them.
|
||||
return Ok(message.trim());
|
||||
}
|
||||
}
|
||||
|
||||
bail!("output processor finished without a Done event")
|
||||
}
|
||||
|
||||
/// Build a decoded-text stream from an already-rendered completion body.
|
||||
///
|
||||
/// The first event carries real prompt token ids so reasoning parsers can
|
||||
/// initialize from the same prompt boundary production uses. Completion text is
|
||||
/// split into small chunks to exercise streaming parser state across marker
|
||||
/// and JSON boundaries.
|
||||
fn decoded_completion_stream(
|
||||
prompt_token_ids: Vec<u32>,
|
||||
completion_body: &str,
|
||||
) -> Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>> {
|
||||
let prompt_token_count = prompt_token_ids.len();
|
||||
let mut events = vec![DecodedTextEvent::Start {
|
||||
prompt_token_ids: Arc::from(prompt_token_ids.into_boxed_slice()),
|
||||
prompt_logprobs: None,
|
||||
}];
|
||||
|
||||
let chunks = split_by_chars(completion_body, 7);
|
||||
if chunks.is_empty() {
|
||||
events.push({
|
||||
DecodedTextEvent::TextDelta {
|
||||
delta: String::new(),
|
||||
token_ids: Vec::new(),
|
||||
logprobs: None,
|
||||
finished: Some(Finished {
|
||||
prompt_token_count: 0,
|
||||
output_token_count: 0,
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
}
|
||||
});
|
||||
} else {
|
||||
let last_index = chunks.len() - 1;
|
||||
for (index, chunk) in chunks.into_iter().enumerate() {
|
||||
let finished = (index == last_index).then(|| Finished {
|
||||
prompt_token_count,
|
||||
output_token_count: completion_body.chars().count(),
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
});
|
||||
events.push(DecodedTextEvent::TextDelta {
|
||||
delta: chunk,
|
||||
token_ids: Vec::new(),
|
||||
logprobs: None,
|
||||
finished,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
stream::iter(events).map(Ok).boxed()
|
||||
}
|
||||
|
||||
/// Split text into chunks containing at most `chunk_chars` Unicode scalar
|
||||
/// values.
|
||||
fn split_by_chars(text: &str, chunk_chars: usize) -> Vec<String> {
|
||||
let mut chunks = Vec::new();
|
||||
let mut start = 0;
|
||||
let mut count = 0;
|
||||
|
||||
for (index, _) in text.char_indices() {
|
||||
if count == chunk_chars {
|
||||
chunks.push(text[start..index].to_string());
|
||||
start = index;
|
||||
count = 0;
|
||||
}
|
||||
count += 1;
|
||||
}
|
||||
|
||||
if start < text.len() {
|
||||
chunks.push(text[start..].to_string());
|
||||
}
|
||||
|
||||
chunks
|
||||
}
|
||||
|
||||
/// Build a chat request fixture with parser-enabling tool-choice semantics.
|
||||
fn roundtrip_request(
|
||||
request_id: impl Into<String>,
|
||||
messages: Vec<ChatMessage>,
|
||||
tools: Vec<ChatTool>,
|
||||
) -> ChatRequest {
|
||||
let mut request = ChatRequest {
|
||||
request_id: request_id.into(),
|
||||
messages,
|
||||
tool_choice: if tools.is_empty() {
|
||||
ChatToolChoice::None
|
||||
} else {
|
||||
ChatToolChoice::Auto
|
||||
},
|
||||
tools,
|
||||
..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());
|
||||
}
|
||||
|
||||
request
|
||||
}
|
||||
|
||||
/// Return the function tools used by the multiple-tool-call fixture.
|
||||
fn test_tools() -> Vec<ChatTool> {
|
||||
vec![
|
||||
ChatTool {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather for a location".to_string()),
|
||||
parameters: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": { "type": "string" }
|
||||
},
|
||||
"required": ["location"]
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
ChatTool {
|
||||
name: "add".to_string(),
|
||||
description: Some("Add two integers".to_string()),
|
||||
parameters: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"y": { "type": "number" },
|
||||
"x": { "type": "number" }
|
||||
},
|
||||
"required": ["y", "x"]
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
]
|
||||
}
|
||||
@@ -17,7 +17,6 @@ anyhow.workspace = true
|
||||
clap.workspace = true
|
||||
educe.workspace = true
|
||||
itertools.workspace = true
|
||||
mimalloc.workspace = true
|
||||
native-tls-vendored = { workspace = true, optional = true }
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
||||
@@ -11,9 +11,6 @@ use vllm_managed_engine::ManagedEngineHandle;
|
||||
|
||||
use crate::cli::{Cli, Command};
|
||||
|
||||
#[global_allocator]
|
||||
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
||||
|
||||
const TOKIO_WORKER_THREADS_ENV: &str = "TOKIO_WORKER_THREADS";
|
||||
const DEFAULT_MAX_TOKIO_WORKER_THREADS: usize = 32;
|
||||
|
||||
|
||||
@@ -595,27 +595,9 @@ impl EngineCoreClient {
|
||||
}
|
||||
|
||||
/// Return whether the engine is currently sleeping at any level.
|
||||
///
|
||||
/// Under data parallel, all engines should agree on the sleep state: a
|
||||
/// divergence signals a control-plane bug. Returns
|
||||
/// `Error::InconsistentUtilityResults` if engines disagree.
|
||||
pub async fn is_sleeping(&self) -> Result<bool> {
|
||||
let results: Vec<bool> = self.call_utility("is_sleeping", ()).await?;
|
||||
// `engine_count >= 1` is enforced during startup handshake, so `results`
|
||||
// is normally non-empty; fall back to a fail-loud error rather than
|
||||
// indexing in case that invariant is ever bypassed.
|
||||
let first = *results.first().ok_or_else(|| Error::InconsistentUtilityResults {
|
||||
method: "is_sleeping".to_string(),
|
||||
values: "[]".to_string(),
|
||||
})?;
|
||||
if results.iter().all(|&v| v == first) {
|
||||
Ok(first)
|
||||
} else {
|
||||
Err(Error::InconsistentUtilityResults {
|
||||
method: "is_sleeping".to_string(),
|
||||
values: format!("{results:?}"),
|
||||
})
|
||||
}
|
||||
// TODO: we only return the result of the first engine here.
|
||||
Ok(self.call_utility("is_sleeping", ()).await?[0])
|
||||
}
|
||||
|
||||
/// Reset the multi-modal cache.
|
||||
@@ -631,30 +613,18 @@ impl EngineCoreClient {
|
||||
}
|
||||
|
||||
/// Reset the prefix cache and optionally the external connector cache.
|
||||
///
|
||||
/// Under data parallel, returns `true` only when every engine confirms the
|
||||
/// reset (AND aggregation).
|
||||
pub async fn reset_prefix_cache(
|
||||
&self,
|
||||
reset_running_requests: bool,
|
||||
reset_connector: bool,
|
||||
) -> Result<bool> {
|
||||
let results: Vec<bool> = self
|
||||
// TODO: we only return the result of the first engine here.
|
||||
Ok(self
|
||||
.call_utility(
|
||||
"reset_prefix_cache",
|
||||
(reset_running_requests, reset_connector),
|
||||
)
|
||||
.await?;
|
||||
// `engine_count >= 1` is enforced during startup handshake, so `results`
|
||||
// is normally non-empty; fail loud rather than reporting a vacuous
|
||||
// success (`[].all() == true`) in case that invariant is ever bypassed.
|
||||
if results.is_empty() {
|
||||
return Err(Error::InconsistentUtilityResults {
|
||||
method: "reset_prefix_cache".to_string(),
|
||||
values: "[]".to_string(),
|
||||
});
|
||||
}
|
||||
Ok(results.into_iter().all(|ok| ok))
|
||||
.await?[0])
|
||||
}
|
||||
|
||||
/// Put the engine to sleep.
|
||||
|
||||
@@ -107,13 +107,13 @@ impl ClientInner {
|
||||
Ok(registry.abortable_request_ids(request_ids))
|
||||
}
|
||||
|
||||
/// Obtain stream senders for a whole engine output batch with one registry
|
||||
/// lock acquisition.
|
||||
pub fn take_senders_for_outputs<'a>(
|
||||
/// Obtain the stream sender for one output. If it indicates the request is
|
||||
/// finished, it will be removed from the registry.
|
||||
pub fn take_sender_for_output(
|
||||
&self,
|
||||
outputs: impl IntoIterator<Item = &'a EngineCoreOutput>,
|
||||
) -> Vec<Option<mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>>> {
|
||||
self.request_reg.lock().senders_for_outputs(outputs)
|
||||
output: &EngineCoreOutput,
|
||||
) -> Option<mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>> {
|
||||
self.request_reg.lock().sender_for_output(output)
|
||||
}
|
||||
|
||||
/// Remove a batch of requests that have finished or aborted, returning
|
||||
@@ -301,10 +301,9 @@ pub(crate) async fn run_output_dispatcher_loop(
|
||||
|
||||
match outputs.classify() {
|
||||
ClassifiedEngineCoreOutputs::RequestBatch(batch) => {
|
||||
let senders = inner.take_senders_for_outputs(&batch.outputs);
|
||||
for (output, sender) in batch.outputs.into_iter().zip(senders) {
|
||||
for output in batch.outputs {
|
||||
let request_id = output.request_id.clone();
|
||||
let Some(sender) = sender else {
|
||||
let Some(sender) = inner.take_sender_for_output(&output) else {
|
||||
debug!(request_id, "dropping output for inactive request");
|
||||
continue;
|
||||
};
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
@@ -80,7 +80,7 @@ impl EngineRoutingState {
|
||||
#[derive(Debug)]
|
||||
pub struct RequestRegistry {
|
||||
closed: bool,
|
||||
requests: HashMap<String, TrackedRequest>,
|
||||
requests: BTreeMap<String, TrackedRequest>,
|
||||
routing_per_engine: BTreeMap<EngineId, EngineRoutingState>,
|
||||
}
|
||||
|
||||
@@ -88,7 +88,7 @@ impl RequestRegistry {
|
||||
pub fn new(engines: &[ConnectedEngine]) -> Self {
|
||||
Self {
|
||||
closed: false,
|
||||
requests: HashMap::default(),
|
||||
requests: BTreeMap::default(),
|
||||
routing_per_engine: engines
|
||||
.iter()
|
||||
.map(|engine| (engine.engine_id.clone(), EngineRoutingState::default()))
|
||||
@@ -180,15 +180,6 @@ impl RequestRegistry {
|
||||
}
|
||||
}
|
||||
|
||||
/// Obtain stream senders for a whole engine output batch under one
|
||||
/// registry lock. Finished outputs are removed before returning.
|
||||
pub fn senders_for_outputs<'a>(
|
||||
&mut self,
|
||||
outputs: impl IntoIterator<Item = &'a EngineCoreOutput>,
|
||||
) -> Vec<Option<OutputSender>> {
|
||||
outputs.into_iter().map(|output| self.sender_for_output(output)).collect()
|
||||
}
|
||||
|
||||
/// Remove a batch of requests that have finished or aborted, returning
|
||||
/// their stream senders.
|
||||
pub fn finish_many<'a>(
|
||||
|
||||
@@ -83,8 +83,6 @@ pub enum Error {
|
||||
},
|
||||
#[error("utility call `{method}` closed unexpectedly (call_id={call_id})")]
|
||||
UtilityCallClosed { method: String, call_id: u64 },
|
||||
#[error("utility call `{method}` returned inconsistent results across engines: {values}")]
|
||||
InconsistentUtilityResults { method: String, values: String },
|
||||
|
||||
/// A special variant to allow cloning the same error.
|
||||
#[error(transparent)]
|
||||
|
||||
@@ -2,7 +2,6 @@ mod client;
|
||||
mod coordinator;
|
||||
mod error;
|
||||
mod metrics;
|
||||
pub mod mock_engine;
|
||||
pub mod protocol;
|
||||
#[cfg(any(test, feature = "test-util"))]
|
||||
pub mod test_utils;
|
||||
|
||||
@@ -1,265 +0,0 @@
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::time::timeout;
|
||||
use zeromq::prelude::{Socket, SocketRecv, SocketSend};
|
||||
use zeromq::util::PeerIdentity;
|
||||
use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage};
|
||||
|
||||
use crate::EngineId;
|
||||
use crate::error::{Error, Result, bail_unexpected_handshake_message};
|
||||
use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
|
||||
use crate::protocol::{ModelDtype, decode_msgpack, encode_msgpack};
|
||||
|
||||
/// Default model length advertised by reusable mock engine helpers.
|
||||
pub const DEFAULT_MOCK_MAX_MODEL_LEN: u64 = 1024 * 1024;
|
||||
/// Default KV block count advertised by reusable mock engine helpers.
|
||||
pub const DEFAULT_MOCK_NUM_GPU_BLOCKS: u64 = 0;
|
||||
|
||||
/// Startup behavior for one mock engine joining a frontend.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MockEngineConfig {
|
||||
/// Whether the engine should advertise itself as local to the frontend.
|
||||
pub local: bool,
|
||||
/// Whether the engine should advertise itself as headless.
|
||||
pub headless: bool,
|
||||
/// Engine-ready payload reported after INIT, including max model length,
|
||||
/// KV block count, and dtype.
|
||||
pub ready_response: EngineCoreReadyResponse,
|
||||
/// Maximum time to wait for IPC endpoints to appear before connecting.
|
||||
pub connect_timeout: Duration,
|
||||
}
|
||||
|
||||
impl Default for MockEngineConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
local: false,
|
||||
headless: true,
|
||||
ready_response: default_ready_response(),
|
||||
connect_timeout: Duration::from_secs(5),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct the ready response used by the standalone mock engine CLI.
|
||||
pub fn default_ready_response() -> EngineCoreReadyResponse {
|
||||
EngineCoreReadyResponse {
|
||||
max_model_len: DEFAULT_MOCK_MAX_MODEL_LEN,
|
||||
num_gpu_blocks: DEFAULT_MOCK_NUM_GPU_BLOCKS,
|
||||
dp_stats_address: None,
|
||||
dtype: Some(ModelDtype::Float32),
|
||||
}
|
||||
}
|
||||
|
||||
/// Coordinator-side sockets used by one mock engine when coordinator mode
|
||||
/// is enabled.
|
||||
pub struct MockCoordinatorSockets {
|
||||
/// Subscription socket that receives coordinator broadcasts such as
|
||||
/// `START_DP_WAVE`.
|
||||
pub input_sub: SubSocket,
|
||||
/// Push socket used to send coordinator-only `EngineCoreOutputs` back to
|
||||
/// the frontend.
|
||||
pub output_push: PushSocket,
|
||||
}
|
||||
|
||||
/// One mock engine's connection to one frontend client.
|
||||
///
|
||||
/// vLLM launches one engine-client pair per API server process. A remote
|
||||
/// engine connects to every advertised input/output pair and uses the request's
|
||||
/// `client_index` to route outputs back to the originating API server.
|
||||
pub struct MockEngineDataSockets {
|
||||
/// Socket used to receive frontend requests.
|
||||
pub dealer: DealerSocket,
|
||||
/// Socket used to publish normal request outputs back to the frontend.
|
||||
pub push: PushSocket,
|
||||
}
|
||||
|
||||
/// Frontend-facing sockets owned by one mock engine.
|
||||
pub struct MockEngineSockets {
|
||||
/// Decoded INIT message sent by the frontend during handshake.
|
||||
pub init: HandshakeInitMessage,
|
||||
/// Data sockets for all frontend clients in client-index order.
|
||||
///
|
||||
/// For Rust frontend this will always be one socket, while for Python frontend
|
||||
/// this may be multiple sockets if there are multiple API server processes.
|
||||
pub data_sockets: Vec<MockEngineDataSockets>,
|
||||
/// Optional coordinator sockets when the client enabled the in-process
|
||||
/// coordinator.
|
||||
pub coordinator: Option<MockCoordinatorSockets>,
|
||||
}
|
||||
|
||||
/// Build a HELLO or READY handshake status payload.
|
||||
fn ready_message(status: &str, config: &MockEngineConfig) -> ReadyMessage {
|
||||
ReadyMessage {
|
||||
status: Some(status.to_string()),
|
||||
local: Some(config.local),
|
||||
headless: Some(config.headless),
|
||||
parallel_config_hash: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert an engine id into a ZMQ DEALER identity.
|
||||
fn peer_identity(engine_id: impl Into<EngineId>) -> Result<PeerIdentity> {
|
||||
let engine_id = engine_id.into();
|
||||
PeerIdentity::try_from(engine_id.clone()).map_err(|error| Error::UnexpectedHandshakeMessage {
|
||||
message: format!(
|
||||
"invalid mock engine identity {:?}: {error}",
|
||||
engine_id.to_vec()
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
/// Wait for an IPC endpoint path to appear before attempting to connect.
|
||||
async fn wait_for_ipc_endpoint(endpoint: &str, connect_timeout: Duration) -> Result<()> {
|
||||
let Some(socket_path) = endpoint.strip_prefix("ipc://") else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
timeout(connect_timeout, async {
|
||||
while !Path::new(socket_path).exists() {
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| Error::HandshakeTimeout {
|
||||
stage: "mock engine IPC endpoint",
|
||||
timeout: connect_timeout,
|
||||
})
|
||||
}
|
||||
|
||||
/// Encode the engine-ready response sent on input socket registration.
|
||||
fn ready_response_payload(config: &MockEngineConfig) -> Result<Vec<u8>> {
|
||||
encode_msgpack(&config.ready_response)
|
||||
}
|
||||
|
||||
/// Join a frontend-owned handshake endpoint and open mock engine sockets.
|
||||
pub async fn connect_to_frontend(
|
||||
engine_handshake: impl AsRef<str>,
|
||||
engine_id: impl Into<EngineId>,
|
||||
config: MockEngineConfig,
|
||||
) -> Result<MockEngineSockets> {
|
||||
let engine_handshake = engine_handshake.as_ref();
|
||||
wait_for_ipc_endpoint(engine_handshake, config.connect_timeout).await?;
|
||||
|
||||
let peer_identity = peer_identity(engine_id)?;
|
||||
let mut options = SocketOptions::default();
|
||||
options.peer_identity(peer_identity.clone());
|
||||
let mut handshake = DealerSocket::with_options(options);
|
||||
handshake.connect(engine_handshake).await?;
|
||||
handshake
|
||||
.send(ZmqMessage::from(encode_msgpack(&ready_message(
|
||||
"HELLO", &config,
|
||||
))?))
|
||||
.await?;
|
||||
|
||||
let init_frames = handshake.recv().await?.into_vec();
|
||||
if init_frames.len() != 1 {
|
||||
bail_unexpected_handshake_message!(
|
||||
"expected one INIT frame from frontend, got {}",
|
||||
init_frames.len()
|
||||
);
|
||||
}
|
||||
let init: HandshakeInitMessage = decode_msgpack(init_frames[0].as_ref())?;
|
||||
|
||||
if init.addresses.inputs.is_empty() {
|
||||
return Err(Error::UnexpectedHandshakeMessage {
|
||||
message: "frontend INIT did not include an input address".to_string(),
|
||||
});
|
||||
}
|
||||
if init.addresses.inputs.len() != init.addresses.outputs.len() {
|
||||
return Err(Error::UnexpectedHandshakeMessage {
|
||||
message: format!(
|
||||
"frontend INIT input/output address count mismatch: {} inputs, {} outputs",
|
||||
init.addresses.inputs.len(),
|
||||
init.addresses.outputs.len()
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
let mut data_sockets = Vec::with_capacity(init.addresses.inputs.len());
|
||||
for (input_address, output_address) in
|
||||
init.addresses.inputs.iter().zip(init.addresses.outputs.iter())
|
||||
{
|
||||
wait_for_ipc_endpoint(input_address, config.connect_timeout).await?;
|
||||
wait_for_ipc_endpoint(output_address, config.connect_timeout).await?;
|
||||
|
||||
let mut input_options = SocketOptions::default();
|
||||
input_options.peer_identity(peer_identity.clone());
|
||||
let mut dealer = DealerSocket::with_options(input_options);
|
||||
dealer.connect(input_address).await?;
|
||||
dealer.send(ZmqMessage::from(ready_response_payload(&config)?)).await?;
|
||||
|
||||
let mut push = PushSocket::new();
|
||||
push.connect(output_address).await?;
|
||||
|
||||
data_sockets.push(MockEngineDataSockets { dealer, push });
|
||||
}
|
||||
|
||||
let coordinator = match (
|
||||
init.addresses.coordinator_input.as_deref(),
|
||||
init.addresses.coordinator_output.as_deref(),
|
||||
) {
|
||||
(Some(coordinator_input), Some(coordinator_output)) => {
|
||||
let mut input_sub = SubSocket::new();
|
||||
input_sub.connect(coordinator_input).await?;
|
||||
input_sub.subscribe("").await?;
|
||||
|
||||
let mut output_push = PushSocket::new();
|
||||
output_push.connect(coordinator_output).await?;
|
||||
|
||||
let ready = input_sub.recv().await?.into_vec();
|
||||
if ready.len() != 1 || ready[0].as_ref() != b"READY" {
|
||||
bail_unexpected_handshake_message!(
|
||||
"expected coordinator READY marker, got {:?}",
|
||||
ready
|
||||
);
|
||||
}
|
||||
|
||||
Some(MockCoordinatorSockets {
|
||||
input_sub,
|
||||
output_push,
|
||||
})
|
||||
}
|
||||
(None, None) => None,
|
||||
_ => bail_unexpected_handshake_message!(
|
||||
"coordinator handshake addresses must be both present or both absent"
|
||||
),
|
||||
};
|
||||
|
||||
handshake
|
||||
.send(ZmqMessage::from(encode_msgpack(&ready_message(
|
||||
"READY", &config,
|
||||
))?))
|
||||
.await?;
|
||||
|
||||
Ok(MockEngineSockets {
|
||||
init,
|
||||
data_sockets,
|
||||
coordinator,
|
||||
})
|
||||
}
|
||||
|
||||
/// Join already-bootstrapped frontend input/output sockets directly.
|
||||
pub async fn connect_to_bootstrapped_frontend(
|
||||
input_address: impl AsRef<str>,
|
||||
output_address: impl AsRef<str>,
|
||||
engine_id: impl Into<EngineId>,
|
||||
config: MockEngineConfig,
|
||||
) -> Result<(DealerSocket, PushSocket)> {
|
||||
let input_address = input_address.as_ref();
|
||||
let output_address = output_address.as_ref();
|
||||
wait_for_ipc_endpoint(input_address, config.connect_timeout).await?;
|
||||
wait_for_ipc_endpoint(output_address, config.connect_timeout).await?;
|
||||
|
||||
let peer_identity = peer_identity(engine_id)?;
|
||||
let mut input_options = SocketOptions::default();
|
||||
input_options.peer_identity(peer_identity);
|
||||
let mut dealer = DealerSocket::with_options(input_options);
|
||||
dealer.connect(input_address).await?;
|
||||
dealer.send(ZmqMessage::from(ready_response_payload(&config)?)).await?;
|
||||
|
||||
let mut push = PushSocket::new();
|
||||
push.connect(output_address).await?;
|
||||
|
||||
Ok((dealer, push))
|
||||
}
|
||||
@@ -36,14 +36,6 @@ fn is_false(v: &bool) -> bool {
|
||||
!v
|
||||
}
|
||||
|
||||
fn default_top_p() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_repetition_penalty() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
mod classified_outputs;
|
||||
pub mod dtype;
|
||||
pub mod handshake;
|
||||
@@ -73,24 +65,6 @@ pub enum EngineCoreRequestType {
|
||||
}
|
||||
|
||||
impl EngineCoreRequestType {
|
||||
/// Decode the single-byte request type frame used on the engine input
|
||||
/// socket. Returns `None` for unrecognized values.
|
||||
pub fn from_frame(frame: &[u8]) -> Option<Self> {
|
||||
let [value] = frame else {
|
||||
return None;
|
||||
};
|
||||
|
||||
match value {
|
||||
0 => Some(Self::Add),
|
||||
1 => Some(Self::Abort),
|
||||
2 => Some(Self::StartDpWave),
|
||||
3 => Some(Self::Utility),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode the request type as the single-byte frame used on the engine
|
||||
/// input socket.
|
||||
pub fn to_frame(self) -> Bytes {
|
||||
Bytes::from_static(match self {
|
||||
Self::Add => b"\x00",
|
||||
@@ -226,17 +200,14 @@ pub struct EngineCoreSamplingParams {
|
||||
/// greedy sampling.
|
||||
pub temperature: f32,
|
||||
/// Cumulative probability threshold for nucleus sampling.
|
||||
#[serde(default = "default_top_p")]
|
||||
pub top_p: f32,
|
||||
/// Maximum number of top tokens to consider. `0` means all tokens.
|
||||
#[serde(default)]
|
||||
pub top_k: u32,
|
||||
/// Random seed used by the sampler when present.
|
||||
pub seed: Option<i64>,
|
||||
/// Maximum number of tokens to generate per output sequence.
|
||||
pub max_tokens: u32,
|
||||
/// Minimum number of tokens to generate before EOS or stop-token handling.
|
||||
#[serde(default)]
|
||||
pub min_tokens: u32,
|
||||
/// Number of log probabilities to return per generated token.
|
||||
///
|
||||
@@ -247,14 +218,12 @@ pub struct EngineCoreSamplingParams {
|
||||
/// `None` disables prompt logprobs. `-1` requests the full vocabulary.
|
||||
pub prompt_logprobs: Option<i32>,
|
||||
/// Minimum probability threshold for token sampling.
|
||||
#[serde(default)]
|
||||
pub min_p: f32,
|
||||
/// Frequency penalty applied by the sampler.
|
||||
pub frequency_penalty: f32,
|
||||
/// Presence penalty applied by the sampler.
|
||||
pub presence_penalty: f32,
|
||||
/// Repetition penalty applied by the sampler.
|
||||
#[serde(default = "default_repetition_penalty")]
|
||||
pub repetition_penalty: f32,
|
||||
/// Token IDs that stop generation.
|
||||
pub stop_token_ids: Vec<u32>,
|
||||
|
||||
@@ -102,7 +102,7 @@ impl<'de> Deserialize<'de> for UtilityCallId {
|
||||
///
|
||||
/// Original Python payload shape:
|
||||
/// `(client_index, call_id, method_name, args)`
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple)]
|
||||
pub struct EngineCoreUtilityRequest {
|
||||
pub client_index: u32,
|
||||
pub call_id: UtilityCallId,
|
||||
|
||||
@@ -1,18 +1,17 @@
|
||||
use std::future::Future;
|
||||
use std::path::Path;
|
||||
use std::pin::Pin;
|
||||
use std::time::Duration;
|
||||
|
||||
use tempfile::TempDir;
|
||||
use tokio::sync::oneshot;
|
||||
use zeromq::{DealerSocket, PushSocket};
|
||||
use zeromq::prelude::{Socket, SocketRecv, SocketSend};
|
||||
use zeromq::util::PeerIdentity;
|
||||
use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage};
|
||||
|
||||
use crate::EngineId;
|
||||
pub use crate::mock_engine::{MockCoordinatorSockets, MockEngineSockets};
|
||||
use crate::mock_engine::{
|
||||
MockEngineConfig, MockEngineDataSockets, connect_to_bootstrapped_frontend, connect_to_frontend,
|
||||
};
|
||||
use crate::protocol::ModelDtype;
|
||||
use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage};
|
||||
use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
|
||||
|
||||
/// Per-test IPC endpoint namespace backed by a unique temporary directory.
|
||||
///
|
||||
@@ -53,29 +52,156 @@ impl IpcNamespace {
|
||||
}
|
||||
}
|
||||
|
||||
fn test_mock_engine_config() -> MockEngineConfig {
|
||||
MockEngineConfig {
|
||||
local: true,
|
||||
headless: true,
|
||||
ready_response: EngineCoreReadyResponse {
|
||||
max_model_len: 4096,
|
||||
num_gpu_blocks: 0,
|
||||
dp_stats_address: None,
|
||||
dtype: Some(ModelDtype::Float32),
|
||||
},
|
||||
..Default::default()
|
||||
/// Construct a standard local READY message used by mock engines in tests.
|
||||
fn ready_message(status: &str) -> ReadyMessage {
|
||||
ReadyMessage {
|
||||
status: Some(status.to_string()),
|
||||
local: Some(true),
|
||||
headless: Some(true),
|
||||
parallel_config_hash: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct a default ready response payload for mock engine input
|
||||
/// registration.
|
||||
fn ready_response_payload() -> Vec<u8> {
|
||||
rmp_serde::to_vec_named(&EngineCoreReadyResponse {
|
||||
max_model_len: 4096,
|
||||
num_gpu_blocks: 0,
|
||||
dp_stats_address: None,
|
||||
dtype: Some(ModelDtype::Float32),
|
||||
})
|
||||
.expect("encode ready response payload")
|
||||
}
|
||||
|
||||
/// Coordinator-side sockets connected by one mock engine when coordinator mode
|
||||
/// is enabled.
|
||||
pub struct MockCoordinatorConnections {
|
||||
/// Subscription socket that receives coordinator broadcasts such as
|
||||
/// `START_DP_WAVE`.
|
||||
pub input_sub: SubSocket,
|
||||
/// Push socket used to send coordinator-only `EngineCoreOutputs` back to
|
||||
/// the frontend.
|
||||
pub output_push: PushSocket,
|
||||
}
|
||||
|
||||
/// Fully connected mock engine transport state used by tests.
|
||||
pub struct MockEngineConnections {
|
||||
/// Decoded INIT message sent by the frontend during handshake.
|
||||
pub init: HandshakeInitMessage,
|
||||
/// Socket used to receive frontend requests.
|
||||
pub dealer: DealerSocket,
|
||||
/// Socket used to publish normal request outputs back to the frontend.
|
||||
pub push: PushSocket,
|
||||
/// Optional coordinator sockets when the client enabled the in-process
|
||||
/// coordinator.
|
||||
pub coordinator: Option<MockCoordinatorConnections>,
|
||||
}
|
||||
|
||||
/// Complete the engine-core handshake and connect mock input/output sockets
|
||||
/// plus optional coordinator sockets.
|
||||
pub async fn setup_mock_engine_sockets(
|
||||
pub async fn setup_mock_engine_connections(
|
||||
engine_handshake: String,
|
||||
engine_id: impl Into<EngineId>,
|
||||
) -> MockEngineSockets {
|
||||
connect_to_frontend(engine_handshake, engine_id, test_mock_engine_config())
|
||||
) -> MockEngineConnections {
|
||||
// Wait for the client to bind the handshake socket before connecting.
|
||||
// A fixed sleep is racy under CI load; instead poll for the socket file.
|
||||
let socket_path = engine_handshake
|
||||
.strip_prefix("ipc://")
|
||||
.expect("handshake address must be ipc://");
|
||||
for _ in 0..100 {
|
||||
if Path::new(socket_path).exists() {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
|
||||
let peer_identity = PeerIdentity::try_from(engine_id.into()).expect("peer id");
|
||||
|
||||
let mut options = SocketOptions::default();
|
||||
options.peer_identity(peer_identity.clone());
|
||||
let mut handshake = DealerSocket::with_options(options);
|
||||
handshake
|
||||
.connect(&engine_handshake)
|
||||
.await
|
||||
.expect("connect mock engine")
|
||||
.expect("connect mock engine handshake socket");
|
||||
handshake
|
||||
.send(ZmqMessage::from(
|
||||
rmp_serde::to_vec_named(&ready_message("HELLO")).expect("encode HELLO ready message"),
|
||||
))
|
||||
.await
|
||||
.expect("send HELLO ready message");
|
||||
|
||||
let init_frames = handshake.recv().await.expect("receive handshake init message").into_vec();
|
||||
assert_eq!(init_frames.len(), 1);
|
||||
let init: HandshakeInitMessage =
|
||||
rmp_serde::from_slice(init_frames[0].as_ref()).expect("decode handshake init message");
|
||||
|
||||
let mut input_options = SocketOptions::default();
|
||||
input_options.peer_identity(peer_identity);
|
||||
let mut dealer = DealerSocket::with_options(input_options);
|
||||
dealer
|
||||
.connect(&init.addresses.inputs[0])
|
||||
.await
|
||||
.expect("connect mock engine input socket");
|
||||
dealer
|
||||
.send(ZmqMessage::from(ready_response_payload()))
|
||||
.await
|
||||
.expect("send mock engine input ready frame");
|
||||
|
||||
let mut push = PushSocket::new();
|
||||
push.connect(&init.addresses.outputs[0])
|
||||
.await
|
||||
.expect("connect mock engine output socket");
|
||||
|
||||
let coordinator = match (
|
||||
init.addresses.coordinator_input.as_deref(),
|
||||
init.addresses.coordinator_output.as_deref(),
|
||||
) {
|
||||
(Some(coordinator_input), Some(coordinator_output)) => {
|
||||
let mut input_sub = SubSocket::new();
|
||||
input_sub
|
||||
.connect(coordinator_input)
|
||||
.await
|
||||
.expect("connect mock engine coordinator input socket");
|
||||
input_sub
|
||||
.subscribe("")
|
||||
.await
|
||||
.expect("subscribe mock engine coordinator input socket");
|
||||
|
||||
let mut output_push = PushSocket::new();
|
||||
output_push
|
||||
.connect(coordinator_output)
|
||||
.await
|
||||
.expect("connect mock engine coordinator output socket");
|
||||
|
||||
let ready =
|
||||
input_sub.recv().await.expect("receive coordinator READY marker").into_vec();
|
||||
assert_eq!(ready.len(), 1);
|
||||
assert_eq!(ready[0].as_ref(), b"READY");
|
||||
|
||||
Some(MockCoordinatorConnections {
|
||||
input_sub,
|
||||
output_push,
|
||||
})
|
||||
}
|
||||
(None, None) => None,
|
||||
_ => panic!("coordinator handshake addresses must be both present or both absent"),
|
||||
};
|
||||
|
||||
handshake
|
||||
.send(ZmqMessage::from(
|
||||
rmp_serde::to_vec_named(&ready_message("READY")).expect("encode READY ready message"),
|
||||
))
|
||||
.await
|
||||
.expect("send READY ready message");
|
||||
|
||||
MockEngineConnections {
|
||||
init,
|
||||
dealer,
|
||||
push,
|
||||
coordinator,
|
||||
}
|
||||
}
|
||||
|
||||
/// Connect one mock engine directly to already-bootstrapped frontend
|
||||
@@ -85,14 +211,31 @@ pub async fn setup_bootstrapped_mock_engine(
|
||||
output_address: String,
|
||||
engine_id: impl Into<EngineId>,
|
||||
) -> (DealerSocket, PushSocket) {
|
||||
connect_to_bootstrapped_frontend(
|
||||
input_address,
|
||||
output_address,
|
||||
engine_id,
|
||||
test_mock_engine_config(),
|
||||
)
|
||||
.await
|
||||
.expect("connect bootstrapped mock engine")
|
||||
for endpoint in [&input_address, &output_address] {
|
||||
if let Some(socket_path) = endpoint.strip_prefix("ipc://") {
|
||||
for _ in 0..100 {
|
||||
if Path::new(socket_path).exists() {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let peer_identity = PeerIdentity::try_from(engine_id.into()).expect("peer id");
|
||||
let mut input_options = SocketOptions::default();
|
||||
input_options.peer_identity(peer_identity);
|
||||
let mut dealer = DealerSocket::with_options(input_options);
|
||||
dealer.connect(&input_address).await.expect("connect mock engine input socket");
|
||||
dealer
|
||||
.send(ZmqMessage::from(ready_response_payload()))
|
||||
.await
|
||||
.expect("send mock engine input ready frame");
|
||||
|
||||
let mut push = PushSocket::new();
|
||||
push.connect(&output_address).await.expect("connect mock engine output socket");
|
||||
|
||||
(dealer, push)
|
||||
}
|
||||
|
||||
/// Complete the engine-core handshake and connect mock input/output sockets.
|
||||
@@ -104,11 +247,9 @@ pub async fn setup_mock_engine_with_init(
|
||||
engine_handshake: String,
|
||||
engine_id: impl Into<EngineId>,
|
||||
) -> (HandshakeInitMessage, DealerSocket, PushSocket) {
|
||||
let MockEngineSockets {
|
||||
init, data_sockets, ..
|
||||
} = setup_mock_engine_sockets(engine_handshake, engine_id).await;
|
||||
let MockEngineDataSockets { dealer, push } =
|
||||
data_sockets.into_iter().next().expect("mock engine data socket");
|
||||
let MockEngineConnections {
|
||||
init, dealer, push, ..
|
||||
} = setup_mock_engine_connections(engine_handshake, engine_id).await;
|
||||
(init, dealer, push)
|
||||
}
|
||||
|
||||
|
||||
@@ -30,7 +30,7 @@ use crate::protocol::{
|
||||
EngineCoreRequestType, EngineCoreSamplingParams, decode_engine_core_outputs,
|
||||
};
|
||||
use crate::test_utils::{
|
||||
IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets,
|
||||
IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_connections,
|
||||
setup_mock_engine_with_init, spawn_mock_engine_task,
|
||||
};
|
||||
use crate::{
|
||||
@@ -477,8 +477,8 @@ async fn coordinator_handshake_includes_engine_control_addresses() {
|
||||
let (init_tx, init_rx) = oneshot::channel();
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel();
|
||||
let engine_task = tokio::spawn(async move {
|
||||
let sockets = setup_mock_engine_sockets(handshake_address, &engine_id).await;
|
||||
let _ = init_tx.send(sockets.init.clone());
|
||||
let connections = setup_mock_engine_connections(handshake_address, &engine_id).await;
|
||||
let _ = init_tx.send(connections.init.clone());
|
||||
let _ = shutdown_rx.await;
|
||||
});
|
||||
|
||||
@@ -515,15 +515,14 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
let engine0_task = tokio::spawn({
|
||||
let handshake_address = handshake_address.clone();
|
||||
async move {
|
||||
let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await;
|
||||
let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
|
||||
let mut coordinator =
|
||||
engine.coordinator.take().expect("coordinator sockets should be present");
|
||||
let data_socket = engine.data_sockets.first_mut().expect("data socket");
|
||||
|
||||
let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
|
||||
assert_eq!((wave, exclude_engine), (0, 0));
|
||||
|
||||
let add = recv_engine_message(&mut data_socket.dealer).await;
|
||||
let add = recv_engine_message(&mut engine.dealer).await;
|
||||
assert_eq!(add[0].as_ref(), &[0x00]);
|
||||
let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
|
||||
assert_eq!(request.request_id, "req-1");
|
||||
@@ -539,7 +538,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
);
|
||||
|
||||
send_outputs(
|
||||
&mut data_socket.push,
|
||||
&mut engine.push,
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![request_output(
|
||||
@@ -566,14 +565,14 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
|
||||
assert_eq!((wave, exclude_engine), (1, 0));
|
||||
|
||||
let add = recv_engine_message(&mut data_socket.dealer).await;
|
||||
let add = recv_engine_message(&mut engine.dealer).await;
|
||||
assert_eq!(add[0].as_ref(), &[0x00]);
|
||||
let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
|
||||
assert_eq!(request.request_id, "req-3");
|
||||
assert_eq!(request.current_wave, 1);
|
||||
|
||||
send_outputs(
|
||||
&mut data_socket.push,
|
||||
&mut engine.push,
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![request_output(
|
||||
@@ -595,15 +594,14 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
let engine1_task = tokio::spawn({
|
||||
let handshake_address = handshake_address.clone();
|
||||
async move {
|
||||
let mut engine = setup_mock_engine_sockets(handshake_address, &[0x01, 0x00]).await;
|
||||
let mut engine = setup_mock_engine_connections(handshake_address, &[0x01, 0x00]).await;
|
||||
let mut coordinator =
|
||||
engine.coordinator.take().expect("coordinator sockets should be present");
|
||||
let data_socket = engine.data_sockets.first_mut().expect("data socket");
|
||||
|
||||
let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
|
||||
assert_eq!((wave, exclude_engine), (0, 0));
|
||||
|
||||
let add = recv_engine_message(&mut data_socket.dealer).await;
|
||||
let add = recv_engine_message(&mut engine.dealer).await;
|
||||
assert_eq!(add[0].as_ref(), &[0x00]);
|
||||
let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
|
||||
assert_eq!(request.request_id, "req-2");
|
||||
@@ -619,7 +617,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
);
|
||||
|
||||
send_outputs(
|
||||
&mut data_socket.push,
|
||||
&mut engine.push,
|
||||
EngineCoreOutputs {
|
||||
engine_index: 1,
|
||||
outputs: vec![request_output(
|
||||
@@ -639,7 +637,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
assert!(
|
||||
timeout(
|
||||
Duration::from_millis(200),
|
||||
recv_engine_message(&mut data_socket.dealer)
|
||||
recv_engine_message(&mut engine.dealer)
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
@@ -714,7 +712,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() {
|
||||
let engine0_task = tokio::spawn({
|
||||
let handshake_address = handshake_address.clone();
|
||||
async move {
|
||||
let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await;
|
||||
let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
|
||||
let mut coordinator =
|
||||
engine.coordinator.take().expect("coordinator sockets should be present");
|
||||
|
||||
@@ -729,7 +727,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() {
|
||||
let engine1_task = tokio::spawn({
|
||||
let handshake_address = handshake_address.clone();
|
||||
async move {
|
||||
let mut engine = setup_mock_engine_sockets(handshake_address, &[0x01, 0x00]).await;
|
||||
let mut engine = setup_mock_engine_connections(handshake_address, &[0x01, 0x00]).await;
|
||||
let mut coordinator =
|
||||
engine.coordinator.take().expect("coordinator sockets should be present");
|
||||
|
||||
@@ -780,10 +778,9 @@ async fn coordinator_accepts_stats_only_outputs() {
|
||||
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel();
|
||||
let engine_task = tokio::spawn(async move {
|
||||
let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await;
|
||||
let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
|
||||
let mut coordinator =
|
||||
engine.coordinator.take().expect("coordinator sockets should be present");
|
||||
let data_socket = engine.data_sockets.first_mut().expect("data socket");
|
||||
|
||||
let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
|
||||
assert_eq!((wave, exclude_engine), (0, 0));
|
||||
@@ -802,13 +799,13 @@ async fn coordinator_accepts_stats_only_outputs() {
|
||||
)
|
||||
.await;
|
||||
|
||||
let add = recv_engine_message(&mut data_socket.dealer).await;
|
||||
let add = recv_engine_message(&mut engine.dealer).await;
|
||||
assert_eq!(add[0].as_ref(), &[0x00]);
|
||||
let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
|
||||
assert_eq!(request.request_id, "req-stats");
|
||||
|
||||
send_outputs(
|
||||
&mut data_socket.push,
|
||||
&mut engine.push,
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![request_output(
|
||||
@@ -2119,226 +2116,6 @@ async fn collective_rpc_flattens_results_from_all_engines() {
|
||||
client.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
/// Spawn a mock engine that handles a single utility call, asserts the method
|
||||
/// name and serialized args match, and replies with `result`.
|
||||
fn spawn_mock_utility_engine(
|
||||
handshake_address: String,
|
||||
engine_id: Vec<u8>,
|
||||
expected_method: &'static str,
|
||||
expected_args: Value,
|
||||
result: bool,
|
||||
) -> (
|
||||
tokio::sync::oneshot::Sender<()>,
|
||||
tokio::task::JoinHandle<()>,
|
||||
) {
|
||||
spawn_mock_engine_task(handshake_address, engine_id, move |dealer, push| {
|
||||
Box::pin(async move {
|
||||
let utility = recv_engine_message(dealer).await;
|
||||
assert_eq!(utility[0].as_ref(), &[0x03]);
|
||||
let payload = decode_value(&utility[1]);
|
||||
let array = match payload {
|
||||
Value::Array(array) => array,
|
||||
other => panic!("expected utility payload array, got {other:?}"),
|
||||
};
|
||||
// Utility requests serialize as `(client_index, call_id, method, args)`.
|
||||
let call_id = array[1].as_u64().expect("call_id");
|
||||
assert_eq!(array[2], Value::from(expected_method));
|
||||
assert_eq!(array[3], expected_args, "unexpected utility args");
|
||||
send_outputs(
|
||||
push,
|
||||
EngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: call_id.into(),
|
||||
failure_message: None,
|
||||
result: Some(utility_result_value(result)),
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn is_sleeping_returns_error_when_engines_disagree() {
|
||||
init_tracing();
|
||||
let ipc = IpcNamespace::new().unwrap();
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
|
||||
let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
|
||||
handshake_address.clone(),
|
||||
b"engine-0".to_vec(),
|
||||
"is_sleeping",
|
||||
Value::Array(vec![]),
|
||||
true,
|
||||
);
|
||||
let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
|
||||
handshake_address.clone(),
|
||||
b"engine-1".to_vec(),
|
||||
"is_sleeping",
|
||||
Value::Array(vec![]),
|
||||
false,
|
||||
);
|
||||
|
||||
let client = connect_client_with_ipc(
|
||||
handshake_test_config(
|
||||
handshake_address,
|
||||
2,
|
||||
"test-model",
|
||||
Duration::from_secs(2),
|
||||
5,
|
||||
None,
|
||||
),
|
||||
&ipc,
|
||||
)
|
||||
.await;
|
||||
|
||||
let error = client.is_sleeping().await.unwrap_err();
|
||||
assert!(
|
||||
matches!(
|
||||
&error,
|
||||
Error::InconsistentUtilityResults { method, .. } if method == "is_sleeping"
|
||||
),
|
||||
"expected InconsistentUtilityResults, got {error:?}",
|
||||
);
|
||||
|
||||
let _ = shutdown_tx_0.send(());
|
||||
let _ = shutdown_tx_1.send(());
|
||||
engine_task_0.await.unwrap();
|
||||
engine_task_1.await.unwrap();
|
||||
client.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn is_sleeping_returns_value_when_all_engines_agree() {
|
||||
init_tracing();
|
||||
let ipc = IpcNamespace::new().unwrap();
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
|
||||
let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
|
||||
handshake_address.clone(),
|
||||
b"engine-0".to_vec(),
|
||||
"is_sleeping",
|
||||
Value::Array(vec![]),
|
||||
true,
|
||||
);
|
||||
let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
|
||||
handshake_address.clone(),
|
||||
b"engine-1".to_vec(),
|
||||
"is_sleeping",
|
||||
Value::Array(vec![]),
|
||||
true,
|
||||
);
|
||||
|
||||
let client = connect_client_with_ipc(
|
||||
handshake_test_config(
|
||||
handshake_address,
|
||||
2,
|
||||
"test-model",
|
||||
Duration::from_secs(2),
|
||||
5,
|
||||
None,
|
||||
),
|
||||
&ipc,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(client.is_sleeping().await.unwrap());
|
||||
|
||||
let _ = shutdown_tx_0.send(());
|
||||
let _ = shutdown_tx_1.send(());
|
||||
engine_task_0.await.unwrap();
|
||||
engine_task_1.await.unwrap();
|
||||
client.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn reset_prefix_cache_returns_true_when_all_engines_succeed() {
|
||||
init_tracing();
|
||||
let ipc = IpcNamespace::new().unwrap();
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
|
||||
let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
|
||||
handshake_address.clone(),
|
||||
b"engine-0".to_vec(),
|
||||
"reset_prefix_cache",
|
||||
Value::Array(vec![Value::from(false), Value::from(false)]),
|
||||
true,
|
||||
);
|
||||
let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
|
||||
handshake_address.clone(),
|
||||
b"engine-1".to_vec(),
|
||||
"reset_prefix_cache",
|
||||
Value::Array(vec![Value::from(false), Value::from(false)]),
|
||||
true,
|
||||
);
|
||||
|
||||
let client = connect_client_with_ipc(
|
||||
handshake_test_config(
|
||||
handshake_address,
|
||||
2,
|
||||
"test-model",
|
||||
Duration::from_secs(2),
|
||||
5,
|
||||
None,
|
||||
),
|
||||
&ipc,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(client.reset_prefix_cache(false, false).await.unwrap());
|
||||
|
||||
let _ = shutdown_tx_0.send(());
|
||||
let _ = shutdown_tx_1.send(());
|
||||
engine_task_0.await.unwrap();
|
||||
engine_task_1.await.unwrap();
|
||||
client.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn reset_prefix_cache_returns_false_when_any_engine_fails() {
|
||||
init_tracing();
|
||||
let ipc = IpcNamespace::new().unwrap();
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
|
||||
let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
|
||||
handshake_address.clone(),
|
||||
b"engine-0".to_vec(),
|
||||
"reset_prefix_cache",
|
||||
Value::Array(vec![Value::from(false), Value::from(false)]),
|
||||
true,
|
||||
);
|
||||
let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
|
||||
handshake_address.clone(),
|
||||
b"engine-1".to_vec(),
|
||||
"reset_prefix_cache",
|
||||
Value::Array(vec![Value::from(false), Value::from(false)]),
|
||||
false,
|
||||
);
|
||||
|
||||
let client = connect_client_with_ipc(
|
||||
handshake_test_config(
|
||||
handshake_address,
|
||||
2,
|
||||
"test-model",
|
||||
Duration::from_secs(2),
|
||||
5,
|
||||
None,
|
||||
),
|
||||
&ipc,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(!client.reset_prefix_cache(false, false).await.unwrap());
|
||||
|
||||
let _ = shutdown_tx_0.send(());
|
||||
let _ = shutdown_tx_1.send(());
|
||||
engine_task_0.await.unwrap();
|
||||
engine_task_1.await.unwrap();
|
||||
client.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn python_msgpack_fixtures_match_rust_encoding() {
|
||||
init_tracing();
|
||||
|
||||
@@ -1,30 +0,0 @@
|
||||
[package]
|
||||
name = "vllm-mock-engine"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[[bin]]
|
||||
name = "vllm-mock-engine"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
anyhow.workspace = true
|
||||
asynk-strim-attr.workspace = true
|
||||
clap.workspace = true
|
||||
futures.workspace = true
|
||||
rand.workspace = true
|
||||
rmpv.workspace = true
|
||||
serde.workspace = true
|
||||
tokio = { workspace = true, features = ["signal"] }
|
||||
tokio-util.workspace = true
|
||||
tracing.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
vllm-engine-core-client.workspace = true
|
||||
zeromq.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
vllm-engine-core-client = { workspace = true, features = ["test-util"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -1,104 +0,0 @@
|
||||
# vLLM Mock Engine
|
||||
|
||||
`vllm-mock-engine` is a small engine-side process for frontend stress testing. It
|
||||
joins a frontend-owned startup handshake, reports a large ready response, treats
|
||||
prefill as instant, and emits random decode tokens until each request reaches
|
||||
its `max_tokens`.
|
||||
|
||||
The frontend must own the handshake socket. Start the frontend first, then start
|
||||
the mock engine with the same handshake address.
|
||||
|
||||
## Start the mock engine
|
||||
|
||||
```bash
|
||||
cargo run -p vllm-mock-engine -- \
|
||||
--handshake-address tcp://127.0.0.1:29550 \
|
||||
--engine-count 1 \
|
||||
--output-token-chunk-size 1 \
|
||||
--vocab-size 32000 \
|
||||
--seed 0 \
|
||||
--log-requests
|
||||
```
|
||||
|
||||
Useful knobs:
|
||||
|
||||
- `--engine-count` must match the frontend's expected data-parallel engine
|
||||
count.
|
||||
- `--output-token-chunk-size` controls how many token IDs appear in one
|
||||
`EngineCoreOutput`; values greater than 1 are useful for MTP/spec-decode
|
||||
shaped frontend tests.
|
||||
- `--vocab-size` should stay within the tokenizer vocabulary of the model used
|
||||
by the frontend.
|
||||
|
||||
Stop it with Ctrl-C.
|
||||
|
||||
## Rust Frontend
|
||||
|
||||
Terminal 1:
|
||||
|
||||
```bash
|
||||
cargo run --bin vllm-rs -- serve Qwen/Qwen3-0.6B \
|
||||
--data-parallel-size 1 \
|
||||
--data-parallel-size-local 0 \
|
||||
--handshake-port 29550
|
||||
```
|
||||
|
||||
Terminal 2:
|
||||
|
||||
```bash
|
||||
cargo run -p vllm-mock-engine -- \
|
||||
--handshake-address tcp://127.0.0.1:29550
|
||||
```
|
||||
|
||||
For multiple mock engines, set both sides to the same count:
|
||||
|
||||
```bash
|
||||
cargo run --bin vllm-rs -- serve Qwen/Qwen3-0.6B \
|
||||
--data-parallel-size 4 \
|
||||
--data-parallel-size-local 0 \
|
||||
--handshake-port 29550
|
||||
|
||||
cargo run -p vllm-mock-engine -- \
|
||||
--handshake-address tcp://127.0.0.1:29550 \
|
||||
--engine-count 4
|
||||
```
|
||||
|
||||
## Python Frontend
|
||||
|
||||
Use `vllm serve` with `--data-parallel-size-local 0` so the Python process runs
|
||||
as a frontend/API server and waits for external engines on
|
||||
`--data-parallel-rpc-port`.
|
||||
|
||||
Terminal 1:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-0.6B \
|
||||
--data-parallel-address 127.0.0.1 \
|
||||
--data-parallel-rpc-port 29550 \
|
||||
--data-parallel-size 1 \
|
||||
--data-parallel-size-local 0
|
||||
```
|
||||
|
||||
Terminal 2:
|
||||
|
||||
```bash
|
||||
cargo run -p vllm-mock-engine -- \
|
||||
--handshake-address tcp://127.0.0.1:29550
|
||||
```
|
||||
|
||||
## Smoke Request
|
||||
|
||||
After either frontend is ready:
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8000/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Qwen/Qwen3-0.6B",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 16,
|
||||
"stream": true
|
||||
}'
|
||||
```
|
||||
|
||||
Always pass `max_tokens`; the mock engine stops by length.
|
||||
@@ -1,416 +0,0 @@
|
||||
use std::collections::{BTreeMap, BTreeSet, HashMap};
|
||||
use std::hash::{Hash as _, Hasher as _};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
use rand::rngs::StdRng;
|
||||
use rand::{Rng as _, SeedableRng as _};
|
||||
use rmpv::Value;
|
||||
use serde::Serialize;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::task::yield_now;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, info, warn};
|
||||
use vllm_engine_core_client::protocol::utility::{
|
||||
EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::{
|
||||
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
|
||||
};
|
||||
|
||||
use super::Opt;
|
||||
|
||||
/// Derive a stable per-request seed from the CLI seed, engine, and request id.
|
||||
fn request_seed(base_seed: u64, engine_index: u32, request_id: &str) -> u64 {
|
||||
let mut hasher = std::hash::DefaultHasher::new();
|
||||
base_seed.hash(&mut hasher);
|
||||
engine_index.hash(&mut hasher);
|
||||
request_id.hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
|
||||
/// Current UNIX timestamp in seconds for engine-core output envelopes.
|
||||
fn now_secs() -> f64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| duration.as_secs_f64())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Build one request output with only token IDs and terminal status populated.
|
||||
fn request_output(
|
||||
request_id: String,
|
||||
new_token_ids: Vec<u32>,
|
||||
finish_reason: Option<EngineCoreFinishReason>,
|
||||
) -> EngineCoreOutput {
|
||||
EngineCoreOutput {
|
||||
request_id,
|
||||
new_token_ids,
|
||||
finish_reason,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Produce an empty output with a terminal finish reason for an invalid request.
|
||||
fn empty_finish_outputs(
|
||||
engine_index: u32,
|
||||
request_id: String,
|
||||
finish_reason: EngineCoreFinishReason,
|
||||
) -> EngineCoreOutputs {
|
||||
let output = request_output(request_id, Vec::new(), Some(finish_reason));
|
||||
let finished_requests = BTreeSet::from([output.request_id.clone()]);
|
||||
|
||||
EngineCoreOutputs {
|
||||
engine_index,
|
||||
outputs: vec![output],
|
||||
timestamp: now_secs(),
|
||||
finished_requests: Some(finished_requests),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode a utility result into the protocol's msgpack value envelope.
|
||||
fn utility_envelope<T>(value: T) -> Result<UtilityResultEnvelope>
|
||||
where
|
||||
T: Serialize,
|
||||
{
|
||||
Ok(UtilityResultEnvelope::without_type_info(
|
||||
rmpv::ext::to_value(value)?,
|
||||
))
|
||||
}
|
||||
|
||||
/// Produce the minimal utility responses needed by the Rust frontend.
|
||||
fn utility_response(
|
||||
engine_index: u32,
|
||||
request: EngineCoreUtilityRequest,
|
||||
) -> Result<EngineCoreOutputs> {
|
||||
let result = match request.method_name.as_str() {
|
||||
"get_supported_tasks" => utility_envelope(vec!["generate"]),
|
||||
"is_sleeping" => utility_envelope(false),
|
||||
"reset_prefix_cache" => utility_envelope(true),
|
||||
"reset_mm_cache"
|
||||
| "reset_encoder_cache"
|
||||
| "profile"
|
||||
| "sleep"
|
||||
| "wake_up"
|
||||
| "execute_dummy_batch" => utility_envelope(()),
|
||||
_ => utility_envelope(Value::Nil),
|
||||
}?;
|
||||
|
||||
Ok(EngineCoreOutputs {
|
||||
engine_index,
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: request.call_id,
|
||||
failure_message: None,
|
||||
result: Some(result),
|
||||
}),
|
||||
timestamp: now_secs(),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
/// Message sent from the frontend to the mock engine task to drive the engine loop.
|
||||
pub(crate) enum EngineInput {
|
||||
Request(Box<EngineCoreRequest>),
|
||||
Abort(Vec<String>),
|
||||
Utility(EngineCoreUtilityRequest),
|
||||
StartDpWave,
|
||||
}
|
||||
|
||||
/// Message sent from the mock engine task to the frontend for one engine output batch.
|
||||
pub(crate) struct EngineOutput {
|
||||
pub client_index: u32,
|
||||
pub outputs: EngineCoreOutputs,
|
||||
}
|
||||
|
||||
/// Per-request decode state owned by one mock engine.
|
||||
#[derive(Debug)]
|
||||
struct ActiveRequest {
|
||||
request_id: String,
|
||||
client_index: u32,
|
||||
prompt_len: usize,
|
||||
max_tokens: usize,
|
||||
generated: usize,
|
||||
rng: StdRng,
|
||||
}
|
||||
|
||||
impl ActiveRequest {
|
||||
/// Create a new active request from an incoming EngineCoreRequest, or return an immediate
|
||||
/// finish reason if the request is invalid.
|
||||
fn new(
|
||||
engine_index: u32,
|
||||
request: Box<EngineCoreRequest>,
|
||||
opt: &Opt,
|
||||
) -> Result<Self, EngineCoreFinishReason> {
|
||||
let request_id = request.request_id;
|
||||
let client_index = request.client_index;
|
||||
let prompt_len = request.prompt_token_ids.as_ref().map(Vec::len).unwrap_or_default();
|
||||
|
||||
let Some(sampling_params) = request.sampling_params else {
|
||||
warn!(
|
||||
request_id,
|
||||
"request has no sampling params; returning engine error"
|
||||
);
|
||||
return Err(EngineCoreFinishReason::Error);
|
||||
};
|
||||
let max_tokens = sampling_params.max_tokens as usize;
|
||||
|
||||
if opt.log_requests {
|
||||
info!(
|
||||
request_id,
|
||||
prompt_len,
|
||||
max_tokens,
|
||||
chunk_size = opt.output_token_chunk_size,
|
||||
"mock request started"
|
||||
);
|
||||
}
|
||||
|
||||
if max_tokens == 0 {
|
||||
return Err(EngineCoreFinishReason::Length);
|
||||
}
|
||||
|
||||
Ok(ActiveRequest {
|
||||
rng: StdRng::seed_from_u64(request_seed(opt.seed, engine_index, &request_id)),
|
||||
request_id,
|
||||
client_index,
|
||||
prompt_len,
|
||||
max_tokens,
|
||||
generated: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// Advance this request by one mock engine step.
|
||||
fn step(&mut self, opt: &Opt) -> EngineCoreOutput {
|
||||
let remaining = self.max_tokens - self.generated;
|
||||
let chunk_len = remaining.min(opt.output_token_chunk_size);
|
||||
let mut new_token_ids = Vec::with_capacity(chunk_len);
|
||||
for _ in 0..chunk_len {
|
||||
new_token_ids.push(self.rng.random_range(0..opt.vocab_size));
|
||||
}
|
||||
self.generated += chunk_len;
|
||||
|
||||
let finished = self.generated >= self.max_tokens;
|
||||
request_output(
|
||||
self.request_id.clone(),
|
||||
new_token_ids,
|
||||
finished.then_some(EngineCoreFinishReason::Length),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Internal state for one mock engine instance, owned by the engine loop task.
|
||||
struct Engine {
|
||||
engine_index: u32,
|
||||
opt: Opt,
|
||||
active_requests: HashMap<String, ActiveRequest>,
|
||||
}
|
||||
|
||||
impl Engine {
|
||||
/// Drain one frontend request message received on the input DEALER socket.
|
||||
fn handle_input(&mut self, input: EngineInput) -> Result<Vec<EngineOutput>> {
|
||||
let mut outputs = Vec::new();
|
||||
|
||||
match input {
|
||||
EngineInput::Request(request) => {
|
||||
let request_id = request.request_id.clone();
|
||||
let client_index = request.client_index;
|
||||
|
||||
if self.active_requests.contains_key(&request_id) {
|
||||
warn!(
|
||||
engine_index = self.engine_index,
|
||||
request_id, "duplicate mock request id"
|
||||
);
|
||||
return Ok(vec![EngineOutput {
|
||||
client_index,
|
||||
outputs: empty_finish_outputs(
|
||||
self.engine_index,
|
||||
request_id,
|
||||
EngineCoreFinishReason::Error,
|
||||
),
|
||||
}]);
|
||||
}
|
||||
|
||||
match ActiveRequest::new(self.engine_index, request, &self.opt) {
|
||||
Ok(request) => {
|
||||
self.active_requests.insert(request_id, request);
|
||||
}
|
||||
Err(finish_reason) => {
|
||||
return Ok(vec![EngineOutput {
|
||||
client_index,
|
||||
outputs: empty_finish_outputs(
|
||||
self.engine_index,
|
||||
request_id,
|
||||
finish_reason,
|
||||
),
|
||||
}]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
EngineInput::Abort(request_ids) => {
|
||||
let mut outputs_by_client =
|
||||
BTreeMap::<u32, (Vec<EngineCoreOutput>, BTreeSet<String>)>::new();
|
||||
for request_id in request_ids {
|
||||
if let Some(request) = self.active_requests.remove(&request_id) {
|
||||
let output = request_output(
|
||||
request_id.clone(),
|
||||
Vec::new(),
|
||||
Some(EngineCoreFinishReason::Abort),
|
||||
);
|
||||
let (outputs, finished_requests) = outputs_by_client
|
||||
.entry(request.client_index)
|
||||
.or_insert_with(|| (Vec::new(), BTreeSet::new()));
|
||||
outputs.push(output);
|
||||
finished_requests.insert(request_id.clone());
|
||||
if self.opt.log_requests {
|
||||
info!(request_id, finish_reason = "abort", "mock request aborted");
|
||||
}
|
||||
}
|
||||
}
|
||||
for (client_index, (client_outputs, finished_requests)) in outputs_by_client {
|
||||
outputs.push({
|
||||
let outputs = EngineCoreOutputs {
|
||||
engine_index: self.engine_index,
|
||||
outputs: client_outputs,
|
||||
timestamp: now_secs(),
|
||||
finished_requests: Some(finished_requests),
|
||||
..Default::default()
|
||||
};
|
||||
EngineOutput {
|
||||
client_index,
|
||||
outputs,
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
EngineInput::Utility(request) => {
|
||||
debug!(
|
||||
engine_index = self.engine_index,
|
||||
call_id = %request.call_id,
|
||||
method = request.method_name,
|
||||
"mock utility request"
|
||||
);
|
||||
let client_index = request.client_index;
|
||||
outputs.push({
|
||||
let outputs = utility_response(self.engine_index, request)?;
|
||||
EngineOutput {
|
||||
client_index,
|
||||
outputs,
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
EngineInput::StartDpWave => {
|
||||
debug!(
|
||||
engine_index = self.engine_index,
|
||||
"ignoring START_DP_WAVE in mock engine"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(outputs)
|
||||
}
|
||||
|
||||
/// Advance active requests once and return one batched engine output.
|
||||
fn step(&mut self) -> Vec<EngineOutput> {
|
||||
if self.active_requests.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut outputs_by_client =
|
||||
BTreeMap::<u32, (Vec<EngineCoreOutput>, BTreeSet<String>)>::new();
|
||||
let mut all_finished_requests = BTreeSet::new();
|
||||
|
||||
for request in self.active_requests.values_mut() {
|
||||
let client_index = request.client_index;
|
||||
let output = request.step(&self.opt);
|
||||
let request_id = request.request_id.clone();
|
||||
let finished = output.finished();
|
||||
if output.finished() {
|
||||
all_finished_requests.insert(request_id.clone());
|
||||
if self.opt.log_requests {
|
||||
info!(
|
||||
request_id,
|
||||
prompt_len = request.prompt_len,
|
||||
output_tokens = request.generated,
|
||||
finish_reason = "length",
|
||||
"mock request finished"
|
||||
);
|
||||
}
|
||||
}
|
||||
let (outputs, finished_requests) = outputs_by_client
|
||||
.entry(client_index)
|
||||
.or_insert_with(|| (Vec::new(), BTreeSet::new()));
|
||||
if finished {
|
||||
finished_requests.insert(request_id.clone());
|
||||
}
|
||||
outputs.push(output);
|
||||
}
|
||||
|
||||
for request_id in &all_finished_requests {
|
||||
self.active_requests.remove(request_id);
|
||||
}
|
||||
|
||||
outputs_by_client
|
||||
.into_iter()
|
||||
.filter_map(|(client_index, (outputs, finished_requests))| {
|
||||
(!outputs.is_empty()).then(|| EngineOutput {
|
||||
client_index,
|
||||
outputs: EngineCoreOutputs {
|
||||
engine_index: self.engine_index,
|
||||
outputs,
|
||||
timestamp: now_secs(),
|
||||
finished_requests: (!finished_requests.is_empty())
|
||||
.then_some(finished_requests),
|
||||
..Default::default()
|
||||
},
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Run the main loop for the mock engine, receiving `EngineInput` from `input_rx`
|
||||
/// and sending `EngineOutput` to `output_tx` until `shutdown` is cancelled.
|
||||
pub(crate) async fn run_engine_loop(
|
||||
engine_index: u32,
|
||||
opt: Opt,
|
||||
mut input_rx: mpsc::UnboundedReceiver<EngineInput>,
|
||||
output_tx: mpsc::Sender<EngineOutput>,
|
||||
shutdown: CancellationToken,
|
||||
) -> Result<()> {
|
||||
let mut engine = Engine {
|
||||
engine_index,
|
||||
opt,
|
||||
active_requests: HashMap::new(),
|
||||
};
|
||||
|
||||
loop {
|
||||
let outputs = tokio::select! {
|
||||
biased;
|
||||
_ = shutdown.cancelled() => break,
|
||||
|
||||
input = input_rx.recv() => {
|
||||
let input = input
|
||||
.ok_or_else(|| anyhow!("mock engine input channel closed"))?;
|
||||
engine.handle_input(input)?
|
||||
}
|
||||
|
||||
// If there are active requests, step them once after yielding to the scheduler to
|
||||
// avoid blocking the engine loop while still making steady progress on request outputs.
|
||||
_ = yield_now(), if !engine.active_requests.is_empty() => {
|
||||
engine.step()
|
||||
}
|
||||
};
|
||||
|
||||
for output in outputs {
|
||||
output_tx
|
||||
.send(output)
|
||||
.await
|
||||
.map_err(|_| anyhow!("mock engine IO task shut down"))?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,121 +0,0 @@
|
||||
use anyhow::{Context as _, Result, anyhow, bail};
|
||||
use futures::{Stream, StreamExt as _, stream};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
use vllm_engine_core_client::mock_engine::MockEngineDataSockets;
|
||||
use vllm_engine_core_client::protocol::utility::EngineCoreUtilityRequest;
|
||||
use vllm_engine_core_client::protocol::{
|
||||
EngineCoreRequest, EngineCoreRequestType, decode_msgpack, encode_msgpack,
|
||||
};
|
||||
use zeromq::{DealerSocket, PushSocket, SocketRecv as _, SocketSend as _, ZmqMessage};
|
||||
|
||||
use crate::engine::{EngineInput, EngineOutput};
|
||||
|
||||
/// Send one engine output batch to the client over the appropriate push socket.
|
||||
async fn send_engine_outputs_to_client(
|
||||
push_sockets: &mut [PushSocket],
|
||||
EngineOutput {
|
||||
client_index,
|
||||
outputs,
|
||||
}: EngineOutput,
|
||||
) -> Result<()> {
|
||||
let message = ZmqMessage::from(encode_msgpack(&outputs)?);
|
||||
push_sockets[client_index as usize].send(message).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Create a stream of `EngineInput` by continuously receiving messages from the given dealer socket
|
||||
/// and decoding them into `EngineInput`.
|
||||
fn dealer_input_stream(dealer: DealerSocket) -> impl Stream<Item = Result<EngineInput>> {
|
||||
stream::unfold(dealer, |mut dealer| async {
|
||||
let input = loop {
|
||||
let message =
|
||||
match dealer.recv().await.context("failed to receive message from dealer socket") {
|
||||
Ok(message) => message,
|
||||
Err(err) => break Err(err),
|
||||
};
|
||||
|
||||
match decode_request(message) {
|
||||
Ok(input) => break Ok(input),
|
||||
Err(err) => {
|
||||
warn!(%err, "failed to decode engine request message; ignoring");
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
Some((input, dealer))
|
||||
})
|
||||
}
|
||||
|
||||
/// Decode a `ZmqMessage` into an `EngineInput`. Returns an error if the message is malformed or
|
||||
/// contains an unknown/unsupported request type.
|
||||
fn decode_request(message: ZmqMessage) -> Result<EngineInput> {
|
||||
let frames = message.into_vec();
|
||||
if frames.is_empty() {
|
||||
bail!("empty engine request message");
|
||||
}
|
||||
if frames.len() != 2 {
|
||||
bail!("invalid frame count for engine request: {}", frames.len());
|
||||
}
|
||||
|
||||
let request_type_frame = frames[0].as_ref();
|
||||
let Some(request_type) = EngineCoreRequestType::from_frame(request_type_frame) else {
|
||||
bail!("unknown engine request type: {:?}", request_type_frame);
|
||||
};
|
||||
|
||||
let input = match request_type {
|
||||
EngineCoreRequestType::Add => {
|
||||
let request: Box<EngineCoreRequest> = decode_msgpack(frames[1].as_ref())?;
|
||||
EngineInput::Request(request)
|
||||
}
|
||||
EngineCoreRequestType::Abort => {
|
||||
let request_ids: Vec<String> = decode_msgpack(frames[1].as_ref())?;
|
||||
EngineInput::Abort(request_ids)
|
||||
}
|
||||
EngineCoreRequestType::Utility => {
|
||||
let request: EngineCoreUtilityRequest = decode_msgpack(frames[1].as_ref())?;
|
||||
EngineInput::Utility(request)
|
||||
}
|
||||
EngineCoreRequestType::StartDpWave => EngineInput::StartDpWave,
|
||||
};
|
||||
|
||||
Ok(input)
|
||||
}
|
||||
|
||||
/// Run the main IO loop for the mock engine, continuously receiving and decoding raw messages from
|
||||
/// the dealer sockets, sending them to the engine loop task via `input_tx`, and receiving
|
||||
/// `EngineOutput` from the engine loop task via `output_rx` and sending them to the client over the
|
||||
/// appropriate push socket, until `shutdown` is cancelled.
|
||||
pub(crate) async fn run_io_loop(
|
||||
data_sockets: Vec<MockEngineDataSockets>,
|
||||
input_tx: mpsc::UnboundedSender<EngineInput>,
|
||||
mut output_rx: mpsc::Receiver<EngineOutput>,
|
||||
shutdown: CancellationToken,
|
||||
) -> Result<()> {
|
||||
let (dealers, mut push_sockets): (Vec<_>, Vec<_>) =
|
||||
data_sockets.into_iter().map(|sockets| (sockets.dealer, sockets.push)).unzip();
|
||||
let mut input_streams =
|
||||
stream::select_all(dealers.into_iter().map(dealer_input_stream).map(Box::pin));
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = shutdown.cancelled() => return Ok(()),
|
||||
|
||||
output = output_rx.recv() => {
|
||||
let output = output
|
||||
.ok_or_else(|| anyhow!("mock engine output channel closed"))?;
|
||||
send_engine_outputs_to_client(&mut push_sockets, output).await?;
|
||||
}
|
||||
|
||||
input = input_streams.next() => {
|
||||
let input = input
|
||||
.ok_or_else(|| anyhow!("mock engine input streams closed"))??;
|
||||
input_tx
|
||||
.send(input)
|
||||
.map_err(|_| anyhow!("mock engine state task shut down"))?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,138 +0,0 @@
|
||||
use anyhow::{Context, Result, bail};
|
||||
use clap::Parser;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::task::JoinSet;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{error, info};
|
||||
use vllm_engine_core_client::EngineId;
|
||||
use vllm_engine_core_client::mock_engine::{
|
||||
MockEngineConfig, MockEngineSockets, connect_to_frontend,
|
||||
};
|
||||
|
||||
pub mod engine;
|
||||
pub mod io;
|
||||
|
||||
/// Standalone engine-core protocol emulator for frontend stress testing.
|
||||
#[derive(Debug, Clone, Parser)]
|
||||
#[command(
|
||||
name = "vllm-mock-engine",
|
||||
about = "Run a mock vLLM headless engine for Rust frontend stress testing."
|
||||
)]
|
||||
pub struct Opt {
|
||||
/// Frontend-owned ZMQ handshake address.
|
||||
#[arg(long, default_value = "tcp://127.0.0.1:29550")]
|
||||
pub handshake_address: String,
|
||||
|
||||
/// Number of mock engine identities to register with the frontend.
|
||||
#[arg(long, default_value_t = 1)]
|
||||
pub engine_count: usize,
|
||||
|
||||
/// Number of accepted output tokens included in each EngineCoreOutput.
|
||||
#[arg(long, default_value_t = 1)]
|
||||
pub output_token_chunk_size: usize,
|
||||
|
||||
/// Random token IDs are sampled uniformly from 0..vocab_size.
|
||||
#[arg(long, default_value_t = 32_000)]
|
||||
pub vocab_size: u32,
|
||||
|
||||
/// Base seed for deterministic random token generation.
|
||||
#[arg(long, default_value_t = 0)]
|
||||
pub seed: u64,
|
||||
|
||||
/// Log a summary line for each request.
|
||||
#[arg(long)]
|
||||
pub log_requests: bool,
|
||||
}
|
||||
|
||||
/// Run one mock engine until shutdown or transport failure.
|
||||
async fn run_engine(engine_index: u32, opt: Opt, shutdown: CancellationToken) -> Result<()> {
|
||||
let MockEngineSockets { data_sockets, .. } = connect_to_frontend(
|
||||
&opt.handshake_address,
|
||||
EngineId::from_engine_index(engine_index),
|
||||
MockEngineConfig::default(),
|
||||
)
|
||||
.await
|
||||
.with_context(|| format!("mock engine {engine_index} failed to connect to frontend"))?;
|
||||
|
||||
info!(engine_index, "mock engine connected to frontend");
|
||||
|
||||
let (input_tx, input_rx) = mpsc::unbounded_channel();
|
||||
let (output_tx, output_rx) = mpsc::channel(64);
|
||||
|
||||
// IO loop: dealer -> input_tx, output_rx -> push
|
||||
let mut io_loop = tokio::spawn(io::run_io_loop(
|
||||
data_sockets,
|
||||
input_tx,
|
||||
output_rx,
|
||||
shutdown.clone(),
|
||||
));
|
||||
// Engine loop: input_rx -> engine logic -> output_tx
|
||||
let mut engine_loop = tokio::spawn(engine::run_engine_loop(
|
||||
engine_index,
|
||||
opt,
|
||||
input_rx,
|
||||
output_tx,
|
||||
shutdown.clone(),
|
||||
));
|
||||
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = shutdown.cancelled() => {
|
||||
io_loop.abort();
|
||||
engine_loop.abort();
|
||||
io_loop.await.ok();
|
||||
engine_loop.await.ok();
|
||||
}
|
||||
|
||||
result = &mut io_loop => {
|
||||
error!(engine_index, "mock engine IO loop exited unexpectedly");
|
||||
engine_loop.abort();
|
||||
engine_loop.await.ok();
|
||||
result??;
|
||||
}
|
||||
result = &mut engine_loop => {
|
||||
error!(engine_index, "mock engine loop exited unexpectedly");
|
||||
io_loop.abort();
|
||||
io_loop.await.ok();
|
||||
result??;
|
||||
}
|
||||
}
|
||||
|
||||
info!(engine_index, "mock engine shut down");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run all requested mock engines until cancellation or one engine task fails.
|
||||
pub async fn run(opt: Opt, shutdown: CancellationToken) -> Result<()> {
|
||||
info!(?opt, "starting mock engine");
|
||||
|
||||
let mut engines = JoinSet::new();
|
||||
for engine_index in 0..opt.engine_count {
|
||||
engines.spawn(run_engine(
|
||||
engine_index as u32,
|
||||
opt.clone(),
|
||||
shutdown.clone(),
|
||||
));
|
||||
}
|
||||
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = shutdown.cancelled() => {
|
||||
engines.abort_all();
|
||||
while engines.join_next().await.is_some() {}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
joined = engines.join_next() => {
|
||||
match joined {
|
||||
Some(Ok(Ok(()))) => bail!("mock engine exited unexpectedly"),
|
||||
Some(Ok(Err(error))) => Err(error),
|
||||
Some(Err(error)) => Err(error).context("mock engine task join failed"),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
@@ -1,38 +0,0 @@
|
||||
use anyhow::{Context, Result};
|
||||
use clap::Parser as _;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{Level, info};
|
||||
use vllm_mock_engine::Opt;
|
||||
|
||||
fn init_tracing() {
|
||||
tracing_subscriber::fmt().with_max_level(Level::INFO).init();
|
||||
}
|
||||
|
||||
/// Create a cancellation token that is triggered by Ctrl-C.
|
||||
fn shutdown_signal() -> CancellationToken {
|
||||
let token = CancellationToken::new();
|
||||
let shutdown = token.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
tokio::signal::ctrl_c().await.expect("failed to install Ctrl-C signal handler");
|
||||
info!("received shutdown signal (Ctrl-C), shutting down...");
|
||||
shutdown.cancel();
|
||||
});
|
||||
|
||||
token
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
init_tracing();
|
||||
let opt = Opt::parse();
|
||||
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.context("failed to build Tokio runtime")?;
|
||||
|
||||
runtime.block_on(async move {
|
||||
let shutdown = shutdown_signal();
|
||||
vllm_mock_engine::run(opt, shutdown).await
|
||||
})
|
||||
}
|
||||
@@ -1,197 +0,0 @@
|
||||
use std::net::TcpListener;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Result;
|
||||
use futures::StreamExt as _;
|
||||
use tokio::time::timeout;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use vllm_engine_core_client::protocol::{
|
||||
EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
|
||||
};
|
||||
use vllm_engine_core_client::test_utils::IpcNamespace;
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
|
||||
|
||||
use crate::{Opt, run};
|
||||
|
||||
fn free_tcp_address() -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").expect("bind free port");
|
||||
let port = listener.local_addr().expect("local addr").port();
|
||||
drop(listener);
|
||||
format!("tcp://127.0.0.1:{port}")
|
||||
}
|
||||
|
||||
fn client_config(handshake_address: String, engine_count: usize) -> EngineCoreClientConfig {
|
||||
EngineCoreClientConfig {
|
||||
transport_mode: TransportMode::HandshakeOwner {
|
||||
handshake_address,
|
||||
advertised_host: "127.0.0.1".to_string(),
|
||||
engine_count,
|
||||
ready_timeout: Duration::from_secs(5),
|
||||
local_input_address: None,
|
||||
local_output_address: None,
|
||||
},
|
||||
coordinator_mode: None,
|
||||
model_name: "mock-model".to_string(),
|
||||
client_index: 0,
|
||||
}
|
||||
}
|
||||
|
||||
async fn connect_with_mock(
|
||||
handshake_address: String,
|
||||
engine_count: usize,
|
||||
output_token_chunk_size: usize,
|
||||
) -> (
|
||||
EngineCoreClient,
|
||||
CancellationToken,
|
||||
tokio::task::JoinHandle<Result<()>>,
|
||||
) {
|
||||
let shutdown = CancellationToken::new();
|
||||
let task = tokio::spawn(run(
|
||||
Opt {
|
||||
handshake_address: handshake_address.clone(),
|
||||
engine_count,
|
||||
output_token_chunk_size,
|
||||
vocab_size: 32_000,
|
||||
seed: 0,
|
||||
log_requests: false,
|
||||
},
|
||||
shutdown.clone(),
|
||||
));
|
||||
|
||||
let client = timeout(
|
||||
Duration::from_secs(5),
|
||||
EngineCoreClient::connect(client_config(handshake_address, engine_count)),
|
||||
)
|
||||
.await
|
||||
.expect("client connect timeout")
|
||||
.expect("connect client");
|
||||
|
||||
(client, shutdown, task)
|
||||
}
|
||||
|
||||
fn sample_request(request_id: &str, max_tokens: u32) -> EngineCoreRequest {
|
||||
EngineCoreRequest {
|
||||
request_id: request_id.to_string(),
|
||||
prompt_token_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params: Some(EngineCoreSamplingParams {
|
||||
max_tokens,
|
||||
..EngineCoreSamplingParams::for_test()
|
||||
}),
|
||||
arrival_time: 0.0,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
async fn shutdown_mock(
|
||||
client: EngineCoreClient,
|
||||
shutdown: CancellationToken,
|
||||
task: tokio::task::JoinHandle<Result<()>>,
|
||||
) {
|
||||
client.shutdown().await.expect("client shutdown");
|
||||
shutdown.cancel();
|
||||
task.await.expect("mock join").expect("mock run");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn mock_engine_connects_over_tcp() {
|
||||
let handshake_address = free_tcp_address();
|
||||
let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
|
||||
assert_eq!(client.engine_count(), 1);
|
||||
assert_eq!(client.engine_identities()[0], &[0, 0]);
|
||||
assert_eq!(client.max_model_len(), Some(1024 * 1024));
|
||||
shutdown_mock(client, shutdown, task).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn mock_engine_connects_over_ipc() {
|
||||
let ipc = IpcNamespace::new().expect("ipc namespace");
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
|
||||
assert_eq!(client.engine_count(), 1);
|
||||
assert_eq!(client.engine_identities()[0], &[0, 0]);
|
||||
shutdown_mock(client, shutdown, task).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn mock_engine_registers_multiple_identities() {
|
||||
let handshake_address = free_tcp_address();
|
||||
let (client, shutdown, task) = connect_with_mock(handshake_address, 2, 1).await;
|
||||
assert_eq!(client.engine_count(), 2);
|
||||
assert_eq!(client.engine_identities(), vec![&[0, 0][..], &[1, 0][..]]);
|
||||
shutdown_mock(client, shutdown, task).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn chunk_size_one_outputs_one_token_per_update() {
|
||||
let handshake_address = free_tcp_address();
|
||||
let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
|
||||
let mut stream = client.call(sample_request("req-1", 3)).await.expect("call");
|
||||
|
||||
let first = stream.next().await.expect("first").expect("first ok");
|
||||
assert_eq!(first.new_token_ids.len(), 1);
|
||||
assert_eq!(first.finish_reason, None);
|
||||
let second = stream.next().await.expect("second").expect("second ok");
|
||||
assert_eq!(second.new_token_ids.len(), 1);
|
||||
assert_eq!(second.finish_reason, None);
|
||||
let third = stream.next().await.expect("third").expect("third ok");
|
||||
assert_eq!(third.new_token_ids.len(), 1);
|
||||
assert_eq!(third.finish_reason, Some(EngineCoreFinishReason::Length));
|
||||
assert!(stream.next().await.is_none());
|
||||
|
||||
shutdown_mock(client, shutdown, task).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn chunk_size_clips_final_output_to_max_tokens() {
|
||||
let handshake_address = free_tcp_address();
|
||||
let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 4).await;
|
||||
let mut stream = client.call(sample_request("req-clip", 6)).await.expect("call");
|
||||
|
||||
let first = stream.next().await.expect("first").expect("first ok");
|
||||
assert_eq!(first.new_token_ids.len(), 4);
|
||||
assert_eq!(first.finish_reason, None);
|
||||
let second = stream.next().await.expect("second").expect("second ok");
|
||||
assert_eq!(second.new_token_ids.len(), 2);
|
||||
assert_eq!(second.finish_reason, Some(EngineCoreFinishReason::Length));
|
||||
assert!(stream.next().await.is_none());
|
||||
|
||||
shutdown_mock(client, shutdown, task).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn abort_cancels_active_request_and_emits_terminal_output() {
|
||||
let handshake_address = free_tcp_address();
|
||||
let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
|
||||
let mut stream = client.call(sample_request("req-abort", 1_000_000)).await.expect("call");
|
||||
let first = stream.next().await.expect("first").expect("first ok");
|
||||
assert_eq!(first.finish_reason, None);
|
||||
|
||||
client.abort(&["req-abort".to_string()]).await.expect("abort");
|
||||
|
||||
loop {
|
||||
let output = timeout(Duration::from_secs(5), stream.next())
|
||||
.await
|
||||
.expect("stream timeout")
|
||||
.expect("terminal output")
|
||||
.expect("output ok");
|
||||
if output.finish_reason.is_some() {
|
||||
assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Abort));
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
shutdown_mock(client, shutdown, task).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn utility_requests_return_minimal_success_responses() {
|
||||
let handshake_address = free_tcp_address();
|
||||
let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
|
||||
|
||||
assert!(!client.is_sleeping().await.expect("is sleeping"));
|
||||
assert!(client.reset_prefix_cache(false, false).await.expect("reset prefix cache"));
|
||||
client.reset_mm_cache().await.expect("reset mm cache");
|
||||
client.reset_encoder_cache().await.expect("reset encoder cache");
|
||||
|
||||
shutdown_mock(client, shutdown, task).await;
|
||||
}
|
||||
@@ -26,7 +26,6 @@ vllm-tokenizer.workspace = true
|
||||
[dev-dependencies]
|
||||
expect-test.workspace = true
|
||||
futures.workspace = true
|
||||
serial_test.workspace = true
|
||||
tempfile.workspace = true
|
||||
tokio.workspace = true
|
||||
vllm-llm = { workspace = true, features = ["test-util"] }
|
||||
|
||||
@@ -401,7 +401,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "too slow for CI and requires network access to Hugging Face"]
|
||||
#[ignore = "requires network access to Hugging Face and downloads the real Kimi K2.5 tokenizer"]
|
||||
async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
|
||||
let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
|
||||
.await
|
||||
|
||||
@@ -235,8 +235,6 @@ fn merge_unique_token_ids(
|
||||
mod tests {
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use serial_test::file_serial;
|
||||
|
||||
use super::*;
|
||||
use crate::backend::hf::HfTextBackend;
|
||||
use crate::backend::{SamplingHints, TextBackend as _};
|
||||
@@ -388,7 +386,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[file_serial(hf_qwen3)]
|
||||
#[ignore = "requires network access to Hugging Face"]
|
||||
async fn lower_text_request_uses_real_qwen_generation_defaults() {
|
||||
let backend = HfTextBackend::from_model("Qwen/Qwen3-0.6B")
|
||||
.await
|
||||
@@ -412,8 +410,12 @@ mod tests {
|
||||
default_top_k: Some(
|
||||
20,
|
||||
),
|
||||
default_min_p: None,
|
||||
default_repetition_penalty: None,
|
||||
default_min_p: Some(
|
||||
0.1,
|
||||
),
|
||||
default_repetition_penalty: Some(
|
||||
1.2,
|
||||
),
|
||||
default_max_tokens: None,
|
||||
max_model_len: Some(
|
||||
40960,
|
||||
@@ -437,10 +439,10 @@ mod tests {
|
||||
min_tokens: 0,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.0,
|
||||
min_p: 0.1,
|
||||
frequency_penalty: 0.0,
|
||||
presence_penalty: 0.0,
|
||||
repetition_penalty: 1.0,
|
||||
repetition_penalty: 1.2,
|
||||
stop_token_ids: [
|
||||
151643,
|
||||
],
|
||||
@@ -451,13 +453,6 @@ mod tests {
|
||||
151643,
|
||||
151645,
|
||||
},
|
||||
logit_bias: None,
|
||||
allowed_token_ids: None,
|
||||
bad_words_token_ids: None,
|
||||
structured_outputs: None,
|
||||
logprob_token_ids: None,
|
||||
skip_reading_prefix_cache: None,
|
||||
extra_args: None,
|
||||
}
|
||||
"#]]
|
||||
.assert_debug_eq(¶ms);
|
||||
|
||||
@@ -8,7 +8,6 @@ license.workspace = true
|
||||
test-util = []
|
||||
|
||||
[dependencies]
|
||||
easy-ext.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
||||
@@ -9,7 +9,6 @@ use utils::feed_parser;
|
||||
|
||||
const CHUNK_CHARS: usize = 7;
|
||||
const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
|
||||
const LONG_TOOL_ARGUMENT_REPEATS: usize = 256;
|
||||
|
||||
fn mixed_fixture() -> String {
|
||||
concat!(
|
||||
@@ -49,24 +48,6 @@ fn long_normal_text_fixture() -> String {
|
||||
line.repeat(LONG_NORMAL_TEXT_REPEATS)
|
||||
}
|
||||
|
||||
fn long_tool_argument_fixture() -> String {
|
||||
let line =
|
||||
"<section><p>Literal } and <tool_call|> marker-shaped text inside content.</p></section>\n";
|
||||
format!(
|
||||
concat!(
|
||||
"I will write the file.\n",
|
||||
"<|tool_call>",
|
||||
"call:write_file{{",
|
||||
"path:<|\"|>index.html<|\"|>,",
|
||||
"content:<|\"|>{}<|\"|>",
|
||||
"}}",
|
||||
"<tool_call|>",
|
||||
"Done."
|
||||
),
|
||||
line.repeat(LONG_TOOL_ARGUMENT_REPEATS)
|
||||
)
|
||||
}
|
||||
|
||||
fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
|
||||
Gemma4ToolParser::create(tools).expect("Gemma4 parser should initialize")
|
||||
}
|
||||
@@ -117,7 +98,6 @@ fn run_stream_group(
|
||||
fn bench_gemma4(c: &mut Criterion) {
|
||||
let tools = test_tools();
|
||||
let mixed_text = mixed_fixture();
|
||||
let long_tool_argument = long_tool_argument_fixture();
|
||||
let long_normal_text = long_normal_text_fixture();
|
||||
|
||||
run_stream_group(
|
||||
@@ -130,16 +110,6 @@ fn bench_gemma4(c: &mut Criterion) {
|
||||
2,
|
||||
);
|
||||
|
||||
run_stream_group(
|
||||
c,
|
||||
"gemma4/long_tool_argument",
|
||||
&tools,
|
||||
&long_tool_argument,
|
||||
CHUNK_CHARS,
|
||||
"I will write the file.\nDone.",
|
||||
1,
|
||||
);
|
||||
|
||||
run_stream_group(
|
||||
c,
|
||||
"gemma4/long_normal_text",
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use super::{DeepSeekDsmlToolParser, DsmlTokens};
|
||||
use crate::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::{Result, Tool, ToolParseResult, ToolParser};
|
||||
|
||||
/// Tool parser for DeepSeek V3.2 models.
|
||||
///
|
||||
@@ -33,6 +33,7 @@ impl DeepSeekV32ToolParser {
|
||||
}
|
||||
|
||||
impl ToolParser for DeepSeekV32ToolParser {
|
||||
/// Create a boxed DeepSeek V3.2 tool parser.
|
||||
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
@@ -40,21 +41,20 @@ impl ToolParser for DeepSeekV32ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
/// Preserve DSML special tokens while decoding.
|
||||
fn preserve_special_tokens(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
self.0.parse_into(chunk, output)
|
||||
/// Push one decoded text chunk through the DSML parser.
|
||||
fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
|
||||
self.0.push(chunk)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
/// Flush buffered text and reset parser state.
|
||||
fn finish(&mut self) -> Result<ToolParseResult> {
|
||||
self.0.finish()
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
self.0.reset()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -63,8 +63,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::DeepSeekV32ToolParser;
|
||||
use crate::ToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
let params = params
|
||||
@@ -84,27 +84,27 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v32_parse_complete_without_tool_call_keeps_text() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
let result = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(result.normal_text, "Hello, world!");
|
||||
assert!(result.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v32_parse_complete_extracts_single_tool_call() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete(&build_tool_call(
|
||||
"get_weather",
|
||||
&[("location", "SF"), ("date", "2024-01-16")],
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(result.normal_text.is_empty());
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({
|
||||
"location": "SF",
|
||||
"date": "2024-01-16"
|
||||
@@ -119,16 +119,16 @@ mod tests {
|
||||
"Thinking... {}",
|
||||
build_tool_call("get_weather", &[("location", "NYC")])
|
||||
);
|
||||
let output = parser.parse_complete(&output).unwrap();
|
||||
let result = parser.parse_complete(&output).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Thinking... ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(result.normal_text, "Thinking... ");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v32_parse_complete_converts_schema_types() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete(
|
||||
"<｜DSML｜function_calls>\n\
|
||||
<｜DSML｜invoke name=\"convert\">\n\
|
||||
@@ -142,9 +142,9 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({
|
||||
"whole": 5.0,
|
||||
"flag": true,
|
||||
@@ -158,7 +158,7 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v32_parse_complete_string_attr_overrides_schema_types() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete(
|
||||
"<｜DSML｜function_calls>\n\
|
||||
<｜DSML｜invoke name=\"convert\">\n\
|
||||
@@ -172,9 +172,9 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({
|
||||
"whole": "5.0",
|
||||
"flag": "true",
|
||||
@@ -188,7 +188,7 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v32_parse_complete_unescapes_literal_closing_tags_in_parameter_value() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete(&build_tool_call(
|
||||
"get_weather",
|
||||
&[
|
||||
@@ -202,7 +202,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({
|
||||
"location": "Hangzhou </｜DSML｜parameter></｜DSML｜invoke></｜DSML｜function_calls>",
|
||||
"date": "2026-05-08",
|
||||
@@ -213,7 +213,7 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v32_streaming_extracts_single_tool_call() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = collect_stream(
|
||||
let result = collect_stream(
|
||||
&mut parser,
|
||||
&[
|
||||
"<｜DSML｜function_calls>\n",
|
||||
@@ -224,11 +224,11 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(result.normal_text.is_empty());
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({ "location": "SF" })
|
||||
);
|
||||
}
|
||||
@@ -236,7 +236,7 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v32_streaming_preserves_prefix_text() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = collect_stream(
|
||||
let result = collect_stream(
|
||||
&mut parser,
|
||||
&[
|
||||
"Thinking... ",
|
||||
@@ -248,23 +248,23 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "Thinking... ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(result.normal_text, "Thinking... ");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v32_streaming_without_tool_call_emits_text_incrementally() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = collect_stream(&mut parser, &["Hello, ", "world!"]);
|
||||
let result = collect_stream(&mut parser, &["Hello, ", "world!"]);
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(result.normal_text, "Hello, world!");
|
||||
assert!(result.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v32_streaming_extracts_multiple_tool_calls_in_order() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = collect_stream(
|
||||
let result = collect_stream(
|
||||
&mut parser,
|
||||
&[&format!(
|
||||
"{}\n{}",
|
||||
@@ -274,17 +274,17 @@ mod tests {
|
||||
)],
|
||||
);
|
||||
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[1].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[1].tool_index, 1);
|
||||
assert_eq!(result.calls.len(), 2);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[1].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[0].tool_index, 0);
|
||||
assert_eq!(result.calls[1].tool_index, 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({ "location": "SF" })
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
|
||||
json!({ "location": "NYC" })
|
||||
);
|
||||
}
|
||||
@@ -294,11 +294,11 @@ mod tests {
|
||||
let text = build_tool_call("get_weather", &[("location", "SF")]);
|
||||
let chunks = split_by_chars(&text, 5);
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({ "location": "SF" })
|
||||
);
|
||||
}
|
||||
@@ -306,7 +306,7 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v32_streaming_handles_bpe_chunked_dsml_opener() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = collect_stream(
|
||||
let result = collect_stream(
|
||||
&mut parser,
|
||||
&[
|
||||
"<｜DSML｜",
|
||||
@@ -333,11 +333,11 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(result.normal_text.is_empty());
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({ "location": "Beijing" })
|
||||
);
|
||||
}
|
||||
@@ -345,12 +345,12 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v32_streaming_truncated_parameter_does_not_leak_eos() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
parser.parse_chunk("<｜DSML｜function_calls>\n").unwrap();
|
||||
parser.parse_chunk("<｜DSML｜invoke name=\"get_weather\">\n").unwrap();
|
||||
parser.push("<｜DSML｜function_calls>\n").unwrap();
|
||||
parser.push("<｜DSML｜invoke name=\"get_weather\">\n").unwrap();
|
||||
parser
|
||||
.parse_chunk("<｜DSML｜parameter name=\"location\" string=\"true\">Tokyo")
|
||||
.push("<｜DSML｜parameter name=\"location\" string=\"true\">Tokyo")
|
||||
.unwrap();
|
||||
parser.parse_chunk("<｜end▁of▁sentence｜>").unwrap();
|
||||
parser.push("<｜end▁of▁sentence｜>").unwrap();
|
||||
|
||||
let error = parser.finish().unwrap_err();
|
||||
assert!(error.to_report_string().contains("incomplete DeepSeek DSML tool call"));
|
||||
@@ -358,7 +358,7 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v32_streaming_drops_eos_after_complete_tool_calls() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = collect_stream(
|
||||
let result = collect_stream(
|
||||
&mut parser,
|
||||
&[
|
||||
"<｜DSML｜function_calls>\n",
|
||||
@@ -369,15 +369,15 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(result.normal_text.is_empty());
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v32_streaming_ignores_text_after_complete_tool_calls() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
let output = collect_stream(
|
||||
let result = collect_stream(
|
||||
&mut parser,
|
||||
&[
|
||||
"<｜DSML｜function_calls>\n",
|
||||
@@ -389,19 +389,17 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert!(result.normal_text.is_empty());
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v32_streaming_does_not_emit_incomplete_invoke() {
|
||||
let mut parser = DeepSeekV32ToolParser::new(&test_tools());
|
||||
parser.parse_chunk("<｜DSML｜function_calls>\n").unwrap();
|
||||
parser.parse_chunk("<｜DSML｜invoke name=\"get_weather\">\n").unwrap();
|
||||
parser.push("<｜DSML｜function_calls>\n").unwrap();
|
||||
parser.push("<｜DSML｜invoke name=\"get_weather\">\n").unwrap();
|
||||
parser
|
||||
.parse_chunk(
|
||||
"<｜DSML｜parameter name=\"location\" string=\"true\">SF</｜DSML｜parameter>\n",
|
||||
)
|
||||
.push("<｜DSML｜parameter name=\"location\" string=\"true\">SF</｜DSML｜parameter>\n")
|
||||
.unwrap();
|
||||
|
||||
let error = parser.finish().unwrap_err();
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use super::{DeepSeekDsmlToolParser, DsmlTokens};
|
||||
use crate::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::{Result, Tool, ToolParseResult, ToolParser};
|
||||
|
||||
/// Tool parser for DeepSeek V4 models.
|
||||
///
|
||||
@@ -36,6 +36,7 @@ impl DeepSeekV4ToolParser {
|
||||
}
|
||||
|
||||
impl ToolParser for DeepSeekV4ToolParser {
|
||||
/// Create a boxed DeepSeek V4 tool parser.
|
||||
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
@@ -43,29 +44,27 @@ impl ToolParser for DeepSeekV4ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
/// Preserve DSML special tokens while decoding.
|
||||
fn preserve_special_tokens(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
self.0.parse_into(chunk, output)
|
||||
/// Push one decoded text chunk through the DSML parser.
|
||||
fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
|
||||
self.0.push(chunk)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
/// Flush buffered text and reset parser state.
|
||||
fn finish(&mut self) -> Result<ToolParseResult> {
|
||||
self.0.finish()
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
self.0.reset()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::DeepSeekV4ToolParser;
|
||||
use crate::ToolParserTestExt as _;
|
||||
use super::{DeepSeekV4ToolParser, ToolParser};
|
||||
use crate::test_utils::{collect_stream, test_tools};
|
||||
|
||||
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
@@ -86,18 +85,18 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v4_parse_complete_reuses_dsml_parser_with_tool_calls_token() {
|
||||
let mut parser = DeepSeekV4ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete(&build_tool_call(
|
||||
"get_weather",
|
||||
&[("location", "SF"), ("date", "2024-01-16")],
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert!(result.normal_text.is_empty());
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({
|
||||
"location": "SF",
|
||||
"date": "2024-01-16"
|
||||
@@ -108,7 +107,7 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v4_streaming_handles_tool_calls_token_split_across_chunks() {
|
||||
let mut parser = DeepSeekV4ToolParser::new(&test_tools());
|
||||
let output = collect_stream(
|
||||
let result = collect_stream(
|
||||
&mut parser,
|
||||
&[
|
||||
"Thinking... ",
|
||||
@@ -123,11 +122,11 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "Thinking... ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.normal_text, "Thinking... ");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({ "location": "Beijing" })
|
||||
);
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@ use winnow::token::{literal, rest, take_until};
|
||||
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{parse_buffered_event, safe_text_len, xml_unescape};
|
||||
use super::{Result, ToolCallDelta, ToolParserOutput};
|
||||
use super::{Result, ToolCallDelta, ToolParseResult};
|
||||
use crate::Tool;
|
||||
|
||||
mod deepseek_v32;
|
||||
@@ -89,10 +89,10 @@ impl DeepSeekDsmlToolParser {
|
||||
}
|
||||
|
||||
/// Apply one parsed DSML event to parser state and output.
|
||||
fn apply_event(&mut self, event: DsmlEvent, output: &mut ToolParserOutput) -> Result<()> {
|
||||
fn apply_event(&mut self, event: DsmlEvent, result: &mut ToolParseResult) -> Result<()> {
|
||||
match event {
|
||||
DsmlEvent::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
result.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
}
|
||||
DsmlEvent::ToolCallsStart => self.mode = DsmlMode::ToolBlock,
|
||||
DsmlEvent::Invoke { name, raw_params } => {
|
||||
@@ -112,7 +112,7 @@ impl DeepSeekDsmlToolParser {
|
||||
let arguments = serde_json::to_string(&arguments)
|
||||
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
|
||||
|
||||
output.calls.push(ToolCallDelta {
|
||||
result.calls.push(ToolCallDelta {
|
||||
tool_index: self.emitted_invoke_count,
|
||||
name: Some(name),
|
||||
arguments,
|
||||
@@ -125,41 +125,46 @@ impl DeepSeekDsmlToolParser {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
/// Reset all streaming state.
|
||||
fn reset(&mut self) {
|
||||
self.buffer.clear();
|
||||
self.mode = DsmlMode::Text;
|
||||
self.emitted_invoke_count = 0;
|
||||
std::mem::take(&mut self.buffer)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
/// Push one decoded text chunk through the DSML parser.
|
||||
fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
|
||||
// Extract tool calls from streaming model output.
|
||||
//
|
||||
// Uses a buffer-until-complete-invoke strategy: text is buffered until
|
||||
// a complete invoke block is available, then parsed and emitted in one
|
||||
// shot.
|
||||
self.buffer.push_str(chunk);
|
||||
let mut result = ToolParseResult::default();
|
||||
|
||||
while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
|
||||
parse_next_dsml_event(input, self.mode, self.tokens)
|
||||
})? {
|
||||
self.apply_event(event, output)?;
|
||||
self.apply_event(event, &mut result)?;
|
||||
self.buffer.drain(..consumed_len);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
let mut output = ToolParserOutput::default();
|
||||
/// Flush buffered text and reset parser state.
|
||||
fn finish(&mut self) -> Result<ToolParseResult> {
|
||||
let mut result = ToolParseResult::default();
|
||||
match self.mode {
|
||||
DsmlMode::Text => output.normal_text.push_str(&self.buffer),
|
||||
DsmlMode::Text => result.normal_text.push_str(&self.buffer),
|
||||
DsmlMode::Done => {}
|
||||
DsmlMode::ToolBlock => {
|
||||
self.reset();
|
||||
return Err(parsing_failed!("incomplete DeepSeek DSML tool call"));
|
||||
}
|
||||
}
|
||||
let _ = self.reset();
|
||||
Ok(output)
|
||||
self.reset();
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
|
||||
use crate::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::{Result, Tool, ToolParseResult, ToolParser};
|
||||
|
||||
/// Tool parser for DeepSeek V3 JSON-fenced tool calls.
|
||||
///
|
||||
@@ -25,6 +25,7 @@ impl DeepSeekV3ToolParser {
|
||||
}
|
||||
|
||||
impl ToolParser for DeepSeekV3ToolParser {
|
||||
/// Create a boxed DeepSeek V3 tool parser.
|
||||
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
@@ -32,17 +33,15 @@ impl ToolParser for DeepSeekV3ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
self.0.parse_into(chunk, output)
|
||||
/// Push one decoded text chunk through the DeepSeek V3 parser.
|
||||
fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
|
||||
self.0.push(chunk)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
/// Flush buffered text and reset parser state.
|
||||
fn finish(&mut self) -> Result<ToolParseResult> {
|
||||
self.0.finish()
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
self.0.reset()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -56,7 +55,7 @@ mod tests {
|
||||
V3_JSON_START,
|
||||
};
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::{ToolParseResult, ToolParser};
|
||||
|
||||
fn v3_tool_call(function_name: &str, arguments: &str) -> String {
|
||||
format!(
|
||||
@@ -71,39 +70,39 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v3_parse_complete_without_tool_call_keeps_text() {
|
||||
let mut parser = DeepSeekV3ToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
let result = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(result.normal_text, "Hello, world!");
|
||||
assert!(result.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v3_parse_complete_extracts_raw_json_arguments() {
|
||||
let mut parser = DeepSeekV3ToolParser::new(&test_tools());
|
||||
let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete(&format!(
|
||||
"Let me check.\n{} trailing text",
|
||||
tool_section(&[v3_tool_call("get_weather", arguments)])
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check.\n");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
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 deepseek_v3_does_not_validate_or_normalize_arguments() {
|
||||
let mut parser = DeepSeekV3ToolParser::new(&test_tools());
|
||||
let arguments = r#"{"location":"Tokyo",}"#;
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete(&tool_section(&[v3_tool_call("get_weather", arguments)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -123,23 +122,23 @@ mod tests {
|
||||
TOOL_CALLS_END,
|
||||
];
|
||||
|
||||
let mut output = ToolParserOutput::default();
|
||||
let mut result = ToolParseResult::default();
|
||||
let mut observed_arguments = Vec::new();
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
let next = parser.push(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
);
|
||||
output.append(next);
|
||||
result.append(next);
|
||||
}
|
||||
output.append(parser.finish().unwrap());
|
||||
result.append(parser.finish().unwrap());
|
||||
|
||||
assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
|
||||
assert_eq!(
|
||||
output.coalesce_calls().calls[0].arguments,
|
||||
result.coalesce_calls().calls[0].arguments,
|
||||
r#"{"location":"Beijing"}"#
|
||||
);
|
||||
}
|
||||
@@ -153,11 +152,11 @@ mod tests {
|
||||
let chunks = split_by_chars(&input, 5);
|
||||
let mut parser = DeepSeekV3ToolParser::new(&test_tools());
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
assert_eq!(result.normal_text, "hello ");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -166,10 +165,10 @@ mod tests {
|
||||
let arguments = format!("{{\"text\":\"literal {V3_ARGUMENT_END} inside\"}}");
|
||||
let input = tool_section(&[v3_tool_call("echo", &arguments)]);
|
||||
|
||||
let output = parser.parse_complete(&input).unwrap();
|
||||
let result = parser.parse_complete(&input).unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -181,10 +180,10 @@ mod tests {
|
||||
let chunks = split_by_chars(&input, 7);
|
||||
let mut parser = DeepSeekV3ToolParser::new(&test_tools());
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
ToolParseResult {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
@@ -204,14 +203,14 @@ mod tests {
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
.assert_debug_eq(&output);
|
||||
.assert_debug_eq(&result);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v3_finish_fails_incomplete_tool_call() {
|
||||
let mut parser = DeepSeekV3ToolParser::new(&test_tools());
|
||||
parser
|
||||
.parse_chunk(&format!(
|
||||
.push(&format!(
|
||||
"{TOOL_CALLS_START}{TOOL_CALL_START}function{TOOL_CALL_SEPARATOR}get_weather{V3_JSON_START}{{\"location\""
|
||||
))
|
||||
.unwrap();
|
||||
@@ -229,7 +228,7 @@ mod tests {
|
||||
"{TOOL_CALLS_START}{TOOL_CALL_START}tool{TOOL_CALL_SEPARATOR}get_weather{V3_JSON_START}{{}}"
|
||||
);
|
||||
|
||||
let error = parser.parse_chunk(&input).unwrap_err();
|
||||
let error = parser.push(&input).unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
|
||||
use crate::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::{Result, Tool, ToolParseResult, ToolParser};
|
||||
|
||||
/// Tool parser for DeepSeek V3.1 raw JSON tool calls.
|
||||
///
|
||||
@@ -21,6 +21,7 @@ impl DeepSeekV31ToolParser {
|
||||
}
|
||||
|
||||
impl ToolParser for DeepSeekV31ToolParser {
|
||||
/// Create a boxed DeepSeek V3.1 tool parser.
|
||||
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
@@ -28,17 +29,15 @@ impl ToolParser for DeepSeekV31ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
self.0.parse_into(chunk, output)
|
||||
/// Push one decoded text chunk through the DeepSeek V3.1 parser.
|
||||
fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
|
||||
self.0.push(chunk)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
/// Flush buffered text and reset parser state.
|
||||
fn finish(&mut self) -> Result<ToolParseResult> {
|
||||
self.0.finish()
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
self.0.reset()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -51,7 +50,7 @@ mod tests {
|
||||
TOOL_CALL_END, TOOL_CALL_SEPARATOR, TOOL_CALL_START, TOOL_CALLS_END, TOOL_CALLS_START,
|
||||
};
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
||||
use crate::{ToolParseResult, ToolParser};
|
||||
|
||||
fn v31_tool_call(function_name: &str, arguments: &str) -> String {
|
||||
format!("{TOOL_CALL_START}{function_name}{TOOL_CALL_SEPARATOR}{arguments}{TOOL_CALL_END}")
|
||||
@@ -64,39 +63,39 @@ mod tests {
|
||||
#[test]
|
||||
fn deepseek_v31_parse_complete_without_tool_call_keeps_text() {
|
||||
let mut parser = DeepSeekV31ToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
let result = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(result.normal_text, "Hello, world!");
|
||||
assert!(result.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v31_parse_complete_extracts_raw_json_arguments() {
|
||||
let mut parser = DeepSeekV31ToolParser::new(&test_tools());
|
||||
let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete(&format!(
|
||||
"Let me check.{} trailing text",
|
||||
tool_section(&[v31_tool_call("get_weather", arguments)])
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check.");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].tool_index, 0);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(result.normal_text, "Let me check.");
|
||||
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 deepseek_v31_does_not_validate_or_normalize_arguments() {
|
||||
let mut parser = DeepSeekV31ToolParser::new(&test_tools());
|
||||
let arguments = r#"{"location":"Tokyo",}"#;
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete(&tool_section(&[v31_tool_call("get_weather", arguments)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -114,23 +113,23 @@ mod tests {
|
||||
TOOL_CALLS_END,
|
||||
];
|
||||
|
||||
let mut output = ToolParserOutput::default();
|
||||
let mut result = ToolParseResult::default();
|
||||
let mut observed_arguments = Vec::new();
|
||||
for chunk in chunks {
|
||||
let next = parser.parse_chunk(chunk).unwrap();
|
||||
let next = parser.push(chunk).unwrap();
|
||||
observed_arguments.extend(
|
||||
next.calls
|
||||
.iter()
|
||||
.filter(|call| call.name.is_none())
|
||||
.map(|call| call.arguments.clone()),
|
||||
);
|
||||
output.append(next);
|
||||
result.append(next);
|
||||
}
|
||||
output.append(parser.finish().unwrap());
|
||||
result.append(parser.finish().unwrap());
|
||||
|
||||
assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
|
||||
assert_eq!(
|
||||
output.coalesce_calls().calls[0].arguments,
|
||||
result.coalesce_calls().calls[0].arguments,
|
||||
r#"{"location":"Beijing"}"#
|
||||
);
|
||||
}
|
||||
@@ -144,11 +143,11 @@ mod tests {
|
||||
let chunks = split_by_chars(&input, 5);
|
||||
let mut parser = DeepSeekV31ToolParser::new(&test_tools());
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
assert_eq!(result.normal_text, "hello ");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -157,10 +156,10 @@ mod tests {
|
||||
let arguments = format!(r#"{{"text":"literal {TOOL_CALL_END} inside"}}"#);
|
||||
let input = tool_section(&[v31_tool_call("echo", &arguments)]);
|
||||
|
||||
let output = parser.parse_complete(&input).unwrap();
|
||||
let result = parser.parse_complete(&input).unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, arguments);
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].arguments, arguments);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -172,10 +171,10 @@ mod tests {
|
||||
let chunks = split_by_chars(&input, 7);
|
||||
let mut parser = DeepSeekV31ToolParser::new(&test_tools());
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
ToolParseResult {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
@@ -195,7 +194,7 @@ mod tests {
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
.assert_debug_eq(&output);
|
||||
.assert_debug_eq(&result);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -206,18 +205,18 @@ mod tests {
|
||||
);
|
||||
let mut parser = DeepSeekV31ToolParser::new(&test_tools());
|
||||
|
||||
let output = collect_stream(&mut parser, &[&input]);
|
||||
let result = collect_stream(&mut parser, &[&input]);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
assert!(result.normal_text.is_empty());
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v31_finish_fails_incomplete_tool_call() {
|
||||
let mut parser = DeepSeekV31ToolParser::new(&test_tools());
|
||||
parser
|
||||
.parse_chunk(&format!(
|
||||
.push(&format!(
|
||||
"{TOOL_CALLS_START}{TOOL_CALL_START}get_weather{TOOL_CALL_SEPARATOR}{{\"location\""
|
||||
))
|
||||
.unwrap();
|
||||
@@ -233,7 +232,7 @@ mod tests {
|
||||
let mut parser = DeepSeekV31ToolParser::new(&test_tools());
|
||||
let input = format!("{TOOL_CALLS_START}{TOOL_CALL_START}{TOOL_CALL_SEPARATOR}{{}}");
|
||||
|
||||
let error = parser.parse_chunk(&input).unwrap_err();
|
||||
let error = parser.push(&input).unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@ use winnow::stream::Partial;
|
||||
use winnow::token::{literal, rest, take_until};
|
||||
|
||||
use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
|
||||
use super::{Result, ToolCallDelta, ToolParserOutput};
|
||||
use super::{Result, ToolCallDelta, ToolParseResult};
|
||||
|
||||
pub(super) const TOOL_CALLS_START: &str = "<｜tool▁calls▁begin｜>";
|
||||
pub(super) const TOOL_CALLS_END: &str = "<｜tool▁calls▁end｜>";
|
||||
@@ -92,11 +92,11 @@ impl DeepSeekJsonToolParser {
|
||||
fn apply_event(
|
||||
&mut self,
|
||||
event: DeepSeekJsonEvent,
|
||||
output: &mut ToolParserOutput,
|
||||
result: &mut ToolParseResult,
|
||||
) -> Result<()> {
|
||||
match event {
|
||||
DeepSeekJsonEvent::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
result.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
}
|
||||
DeepSeekJsonEvent::ToolCallsStart => self.mode = DeepSeekJsonMode::ToolBlock,
|
||||
DeepSeekJsonEvent::ToolCallStart => self.mode = DeepSeekJsonMode::Header,
|
||||
@@ -107,7 +107,7 @@ impl DeepSeekJsonToolParser {
|
||||
self.mode = DeepSeekJsonMode::Arguments {
|
||||
json_scan: JsonObjectScanState::default(),
|
||||
};
|
||||
output.calls.push(ToolCallDelta {
|
||||
result.calls.push(ToolCallDelta {
|
||||
tool_index,
|
||||
name: Some(function_name),
|
||||
arguments: String::new(),
|
||||
@@ -120,7 +120,7 @@ impl DeepSeekJsonToolParser {
|
||||
self.format.parser_name()
|
||||
));
|
||||
};
|
||||
output.calls.push(ToolCallDelta {
|
||||
result.calls.push(ToolCallDelta {
|
||||
tool_index,
|
||||
name: None,
|
||||
arguments: self.buffer[..consumed_len].to_string(),
|
||||
@@ -139,23 +139,26 @@ impl DeepSeekJsonToolParser {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
/// Push one decoded text chunk through the DeepSeek JSON parser.
|
||||
fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
|
||||
self.buffer.push_str(chunk);
|
||||
let mut result = ToolParseResult::default();
|
||||
|
||||
while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
|
||||
parse_next_deepseek_json_event(input, &mut self.mode, self.format)
|
||||
})? {
|
||||
self.apply_event(event, output)?;
|
||||
self.apply_event(event, &mut result)?;
|
||||
self.buffer.drain(..consumed_len);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
let mut output = ToolParserOutput::default();
|
||||
/// Flush buffered text and reset parser state.
|
||||
fn finish(&mut self) -> Result<ToolParseResult> {
|
||||
let mut result = ToolParseResult::default();
|
||||
match &self.mode {
|
||||
DeepSeekJsonMode::Text => output.normal_text.push_str(&self.buffer),
|
||||
DeepSeekJsonMode::Text => result.normal_text.push_str(&self.buffer),
|
||||
DeepSeekJsonMode::ToolBlock | DeepSeekJsonMode::Done => {}
|
||||
DeepSeekJsonMode::Header | DeepSeekJsonMode::Arguments { .. } => {
|
||||
return Err(parsing_failed!(
|
||||
@@ -164,15 +167,16 @@ impl DeepSeekJsonToolParser {
|
||||
));
|
||||
}
|
||||
}
|
||||
let _ = self.reset();
|
||||
Ok(output)
|
||||
self.reset();
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
/// Reset all streaming state.
|
||||
fn reset(&mut self) {
|
||||
self.buffer.clear();
|
||||
self.mode = DeepSeekJsonMode::Text;
|
||||
self.active_tool_index = None;
|
||||
self.emitted_tool_count = 0;
|
||||
std::mem::take(&mut self.buffer)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+123
-296
@@ -1,13 +1,13 @@
|
||||
use serde_json::{Map, Number, Value};
|
||||
use winnow::ascii::multispace0 as ws0;
|
||||
use winnow::combinator::{alt, delimited, eof, opt, separated, seq, terminated};
|
||||
use winnow::combinator::{alt, delimited, opt, separated, seq, terminated};
|
||||
use winnow::error::{ContextError, ErrMode, ModalResult};
|
||||
use winnow::prelude::*;
|
||||
use winnow::stream::{Partial, Stream};
|
||||
use winnow::stream::Partial;
|
||||
use winnow::token::{literal, take_till, take_until};
|
||||
|
||||
use super::utils::{incomplete, parse_buffered_event, partial_prefix_len, safe_text_len};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use super::utils::{parse_buffered_event, safe_text_len};
|
||||
use super::{Result, ToolCallDelta, ToolParseResult, ToolParser};
|
||||
use crate::Tool;
|
||||
|
||||
const TOOL_CALL_START: &str = "<|tool_call>";
|
||||
@@ -19,26 +19,12 @@ type Gemma4Input<'i> = Partial<&'i str>;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
enum Gemma4Event {
|
||||
Text { len: usize },
|
||||
ToolCallStart,
|
||||
ToolCallHeader { name: String },
|
||||
ToolCall { args: Map<String, Value> },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
struct Gemma4ArgsScanState {
|
||||
scanned_len: usize,
|
||||
in_string: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
enum Gemma4Mode {
|
||||
#[default]
|
||||
Text,
|
||||
Header,
|
||||
Text {
|
||||
len: usize,
|
||||
},
|
||||
ToolCall {
|
||||
name: String,
|
||||
args_scan: Gemma4ArgsScanState,
|
||||
args: Map<String, Value>,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -54,7 +40,6 @@ enum Gemma4Mode {
|
||||
/// Arguments are emitted only after a full Gemma4 tool call is parsed.
|
||||
pub struct Gemma4ToolParser {
|
||||
buffer: String,
|
||||
mode: Gemma4Mode,
|
||||
emitted_tool_count: usize,
|
||||
}
|
||||
|
||||
@@ -62,34 +47,20 @@ impl Gemma4ToolParser {
|
||||
fn new(_tools: &[Tool]) -> Self {
|
||||
Self {
|
||||
buffer: String::new(),
|
||||
mode: Gemma4Mode::default(),
|
||||
emitted_tool_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_event(&mut self, event: Gemma4Event, output: &mut ToolParserOutput) -> Result<()> {
|
||||
fn apply_event(&mut self, event: Gemma4Event, result: &mut ToolParseResult) -> Result<()> {
|
||||
match event {
|
||||
Gemma4Event::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
result.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
}
|
||||
Gemma4Event::ToolCallStart => self.mode = Gemma4Mode::Header,
|
||||
Gemma4Event::ToolCallHeader { name } => {
|
||||
self.mode = Gemma4Mode::ToolCall {
|
||||
name,
|
||||
args_scan: Gemma4ArgsScanState::default(),
|
||||
};
|
||||
}
|
||||
Gemma4Event::ToolCall { args } => {
|
||||
let mode = std::mem::replace(&mut self.mode, Gemma4Mode::Text);
|
||||
let Gemma4Mode::ToolCall { name, .. } = mode else {
|
||||
return Err(parsing_failed!(
|
||||
"Gemma4 arguments without an active tool call"
|
||||
));
|
||||
};
|
||||
Gemma4Event::ToolCall { name, args } => {
|
||||
let arguments = serde_json::to_string(&args)
|
||||
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
|
||||
|
||||
output.calls.push(ToolCallDelta {
|
||||
result.calls.push(ToolCallDelta {
|
||||
tool_index: self.emitted_tool_count,
|
||||
name: Some(name),
|
||||
arguments,
|
||||
@@ -100,25 +71,9 @@ impl Gemma4ToolParser {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
let raw = match std::mem::replace(&mut self.mode, Gemma4Mode::Text) {
|
||||
Gemma4Mode::Text => std::mem::take(&mut self.buffer),
|
||||
Gemma4Mode::Header => {
|
||||
format!("{}{}", TOOL_CALL_START, std::mem::take(&mut self.buffer))
|
||||
}
|
||||
Gemma4Mode::ToolCall { name, .. } => {
|
||||
format!(
|
||||
"{}{}{}{{{}",
|
||||
TOOL_CALL_START,
|
||||
CALL_PREFIX,
|
||||
name,
|
||||
std::mem::take(&mut self.buffer)
|
||||
)
|
||||
}
|
||||
};
|
||||
self.mode = Gemma4Mode::Text;
|
||||
fn reset(&mut self) {
|
||||
self.buffer.clear();
|
||||
self.emitted_tool_count = 0;
|
||||
raw
|
||||
}
|
||||
}
|
||||
|
||||
@@ -134,85 +89,56 @@ impl ToolParser for Gemma4ToolParser {
|
||||
true
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
|
||||
self.buffer.push_str(chunk);
|
||||
let mut result = ToolParseResult::default();
|
||||
|
||||
while let Some((event, consumed_len)) = {
|
||||
parse_buffered_event(&self.buffer, |input| {
|
||||
parse_next_gemma4_event(input, &mut self.mode)
|
||||
})?
|
||||
} {
|
||||
self.apply_event(event, output)?;
|
||||
while let Some((event, consumed_len)) =
|
||||
parse_buffered_event(&self.buffer, parse_next_gemma4_event)?
|
||||
{
|
||||
self.apply_event(event, &mut result)?;
|
||||
self.buffer.drain(..consumed_len);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
let mut output = ToolParserOutput::default();
|
||||
fn finish(&mut self) -> Result<ToolParseResult> {
|
||||
let mut result = ToolParseResult::default();
|
||||
|
||||
match &self.mode {
|
||||
Gemma4Mode::Text => output.normal_text.push_str(&self.buffer),
|
||||
Gemma4Mode::Header | Gemma4Mode::ToolCall { .. } => {
|
||||
if !self.buffer.is_empty() {
|
||||
if self.buffer.starts_with(TOOL_CALL_START) {
|
||||
self.reset();
|
||||
return Err(parsing_failed!("incomplete Gemma4 tool call"));
|
||||
}
|
||||
result.normal_text.push_str(&self.buffer);
|
||||
}
|
||||
|
||||
let _ = self.reset();
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
Gemma4ToolParser::reset(self)
|
||||
self.reset();
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse one Gemma4 event from buffered streaming input.
|
||||
fn parse_next_gemma4_event(
|
||||
input: &mut Gemma4Input<'_>,
|
||||
mode: &mut Gemma4Mode,
|
||||
) -> ModalResult<Gemma4Event> {
|
||||
match mode {
|
||||
Gemma4Mode::Text => parse_text_event(input),
|
||||
Gemma4Mode::Header => tool_call_header_event(input),
|
||||
Gemma4Mode::ToolCall { args_scan, .. } => tool_call_args_event(input, args_scan),
|
||||
}
|
||||
fn parse_next_gemma4_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
|
||||
alt((tool_call_event, safe_text_event)).parse_next(input)
|
||||
}
|
||||
|
||||
/// Parse a Gemma4 text-mode event.
|
||||
fn parse_text_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
|
||||
alt((tool_call_start_event, safe_text_event)).parse_next(input)
|
||||
}
|
||||
|
||||
/// Parse a Gemma4 tool-call start marker.
|
||||
fn tool_call_start_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
|
||||
literal(TOOL_CALL_START).value(Gemma4Event::ToolCallStart).parse_next(input)
|
||||
}
|
||||
|
||||
/// Parse a Gemma4 tool-call header.
|
||||
fn tool_call_header_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
|
||||
let (name,) = seq!(
|
||||
/// Parse a complete Gemma4 tool call.
|
||||
// TODO: incremental parsing arguments to reduce scanning from O(n^2) to O(n).
|
||||
fn tool_call_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
|
||||
let (name, args) = seq!(
|
||||
_: literal(TOOL_CALL_START),
|
||||
_: literal(CALL_PREFIX),
|
||||
gemma4_tool_name,
|
||||
_: literal("{"),
|
||||
gemma4_args,
|
||||
_: literal("}"),
|
||||
_: literal(TOOL_CALL_END),
|
||||
)
|
||||
.parse_next(input)?;
|
||||
Ok(Gemma4Event::ToolCallHeader { name })
|
||||
}
|
||||
|
||||
/// Parse complete Gemma4 tool-call arguments.
|
||||
fn tool_call_args_event(
|
||||
input: &mut Gemma4Input<'_>,
|
||||
args_scan: &mut Gemma4ArgsScanState,
|
||||
) -> ModalResult<Gemma4Event> {
|
||||
let raw_args = gemma4_raw_args_until_tool_call_end(input, args_scan)?;
|
||||
let Some(args_input) = raw_args.strip_suffix('}') else {
|
||||
return Err(ErrMode::Cut(ContextError::new()));
|
||||
};
|
||||
let args = parse_gemma4_args(args_input)?;
|
||||
|
||||
Ok(Gemma4Event::ToolCall { args })
|
||||
Ok(Gemma4Event::ToolCall { name, args })
|
||||
}
|
||||
|
||||
/// Parse a Gemma4 tool name.
|
||||
@@ -229,75 +155,8 @@ fn safe_text_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
|
||||
safe_text_len(input, TOOL_CALL_START).map(|len| Gemma4Event::Text { len })
|
||||
}
|
||||
|
||||
/// Parse raw Gemma4 arguments through the first end marker outside a Gemma string.
|
||||
fn gemma4_raw_args_until_tool_call_end<'i>(
|
||||
input: &mut Gemma4Input<'i>,
|
||||
state: &mut Gemma4ArgsScanState,
|
||||
) -> ModalResult<&'i str> {
|
||||
let text = **input;
|
||||
if state.scanned_len > text.len() {
|
||||
return incomplete();
|
||||
}
|
||||
|
||||
loop {
|
||||
let rest = &text[state.scanned_len..];
|
||||
if state.in_string {
|
||||
let Some(string_delim) = rest.find(STRING_DELIM) else {
|
||||
state.scanned_len = safe_scan_len(text, state.scanned_len, &[STRING_DELIM]);
|
||||
return incomplete();
|
||||
};
|
||||
|
||||
state.scanned_len += string_delim + STRING_DELIM.len();
|
||||
state.in_string = false;
|
||||
continue;
|
||||
}
|
||||
|
||||
let next_string_delim = rest.find(STRING_DELIM);
|
||||
let next_tool_call_end = rest.find(TOOL_CALL_END);
|
||||
match (next_string_delim, next_tool_call_end) {
|
||||
(Some(string_delim), Some(tool_call_end)) if tool_call_end < string_delim => {
|
||||
let end = state.scanned_len + tool_call_end;
|
||||
state.scanned_len = end + TOOL_CALL_END.len();
|
||||
input.next_slice(state.scanned_len);
|
||||
return Ok(&text[..end]);
|
||||
}
|
||||
(Some(string_delim), _) => {
|
||||
state.scanned_len += string_delim + STRING_DELIM.len();
|
||||
state.in_string = true;
|
||||
}
|
||||
(None, Some(tool_call_end)) => {
|
||||
let end = state.scanned_len + tool_call_end;
|
||||
state.scanned_len = end + TOOL_CALL_END.len();
|
||||
input.next_slice(state.scanned_len);
|
||||
return Ok(&text[..end]);
|
||||
}
|
||||
(None, None) => {
|
||||
state.scanned_len =
|
||||
safe_scan_len(text, state.scanned_len, &[STRING_DELIM, TOOL_CALL_END]);
|
||||
return incomplete();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the scan length while holding back a split marker prefix.
|
||||
fn safe_scan_len(text: &str, start: usize, markers: &[&str]) -> usize {
|
||||
let max_partial = markers
|
||||
.iter()
|
||||
.map(|marker| partial_prefix_len(&text[start..], marker))
|
||||
.max()
|
||||
.unwrap_or(0);
|
||||
text.len() - max_partial
|
||||
}
|
||||
|
||||
/// Parse complete Gemma4 custom key-value arguments.
|
||||
fn parse_gemma4_args(args: &str) -> ModalResult<Map<String, Value>> {
|
||||
let mut input = args;
|
||||
terminated(gemma4_args, eof).parse_next(&mut input)
|
||||
}
|
||||
|
||||
/// Parse Gemma4's custom key-value argument object content.
|
||||
fn gemma4_args(input: &mut &str) -> ModalResult<Map<String, Value>> {
|
||||
fn gemma4_args(input: &mut Gemma4Input<'_>) -> ModalResult<Map<String, Value>> {
|
||||
let pairs: Vec<(String, Value)> = delimited(
|
||||
ws0,
|
||||
terminated(
|
||||
@@ -311,7 +170,7 @@ fn gemma4_args(input: &mut &str) -> ModalResult<Map<String, Value>> {
|
||||
}
|
||||
|
||||
/// Parse a Gemma4 key-value pair.
|
||||
fn gemma4_pair(input: &mut &str) -> ModalResult<(String, Value)> {
|
||||
fn gemma4_pair(input: &mut Gemma4Input<'_>) -> ModalResult<(String, Value)> {
|
||||
let (key, value) = seq!(
|
||||
_: ws0,
|
||||
gemma4_key,
|
||||
@@ -325,7 +184,7 @@ fn gemma4_pair(input: &mut &str) -> ModalResult<(String, Value)> {
|
||||
}
|
||||
|
||||
/// Parse a Gemma4 bare key.
|
||||
fn gemma4_key(input: &mut &str) -> ModalResult<String> {
|
||||
fn gemma4_key(input: &mut Gemma4Input<'_>) -> ModalResult<String> {
|
||||
let key = take_till(1.., |char: char| char == ':').parse_next(input)?.trim();
|
||||
if key.is_empty() {
|
||||
return Err(ErrMode::Cut(ContextError::new()));
|
||||
@@ -334,7 +193,7 @@ fn gemma4_key(input: &mut &str) -> ModalResult<String> {
|
||||
}
|
||||
|
||||
/// Parse a Gemma4 value.
|
||||
fn gemma4_value(input: &mut &str) -> ModalResult<Value> {
|
||||
fn gemma4_value(input: &mut Gemma4Input<'_>) -> ModalResult<Value> {
|
||||
alt((
|
||||
gemma4_string.map(|value: &str| Value::String(value.to_string())),
|
||||
gemma4_object.map(Value::Object),
|
||||
@@ -345,7 +204,7 @@ fn gemma4_value(input: &mut &str) -> ModalResult<Value> {
|
||||
}
|
||||
|
||||
/// Parse a Gemma4 string delimited by `<|"|>`.
|
||||
fn gemma4_string<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
|
||||
fn gemma4_string<'i>(input: &mut Gemma4Input<'i>) -> ModalResult<&'i str> {
|
||||
delimited(
|
||||
literal(STRING_DELIM),
|
||||
take_until(0.., STRING_DELIM),
|
||||
@@ -355,17 +214,17 @@ fn gemma4_string<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
|
||||
}
|
||||
|
||||
/// Parse a nested Gemma4 object.
|
||||
fn gemma4_object(input: &mut &str) -> ModalResult<Map<String, Value>> {
|
||||
fn gemma4_object(input: &mut Gemma4Input<'_>) -> ModalResult<Map<String, Value>> {
|
||||
delimited(literal("{"), gemma4_args, literal("}")).parse_next(input)
|
||||
}
|
||||
|
||||
/// Parse a Gemma4 array value.
|
||||
fn gemma4_array_value(input: &mut &str) -> ModalResult<Vec<Value>> {
|
||||
fn gemma4_array_value(input: &mut Gemma4Input<'_>) -> ModalResult<Vec<Value>> {
|
||||
delimited(literal("["), gemma4_array_content, literal("]")).parse_next(input)
|
||||
}
|
||||
|
||||
/// Parse Gemma4 array content.
|
||||
fn gemma4_array_content(input: &mut &str) -> ModalResult<Vec<Value>> {
|
||||
fn gemma4_array_content(input: &mut Gemma4Input<'_>) -> ModalResult<Vec<Value>> {
|
||||
delimited(
|
||||
ws0,
|
||||
terminated(
|
||||
@@ -378,14 +237,14 @@ fn gemma4_array_content(input: &mut &str) -> ModalResult<Vec<Value>> {
|
||||
}
|
||||
|
||||
/// Parse a Gemma4 bare scalar.
|
||||
fn gemma4_bare_value(input: &mut &str) -> ModalResult<Value> {
|
||||
fn gemma4_bare_value(input: &mut Gemma4Input<'_>) -> ModalResult<Value> {
|
||||
take_till(1.., |char: char| matches!(char, ',' | '}' | ']'))
|
||||
.map(parse_gemma4_scalar)
|
||||
.parse_next(input)
|
||||
}
|
||||
|
||||
/// Parse a Gemma4 comma separator.
|
||||
fn comma_separator(input: &mut &str) -> ModalResult<()> {
|
||||
fn comma_separator(input: &mut Gemma4Input<'_>) -> ModalResult<()> {
|
||||
delimited(ws0, literal(","), ws0).void().parse_next(input)
|
||||
}
|
||||
|
||||
@@ -423,15 +282,29 @@ mod tests {
|
||||
use winnow::combinator::{eof, terminated};
|
||||
use winnow::error::ErrMode;
|
||||
use winnow::prelude::*;
|
||||
use winnow::stream::Partial;
|
||||
|
||||
use super::{
|
||||
Gemma4ToolParser, ToolCallDelta, ToolParser, ToolParserOutput, gemma4_array_content,
|
||||
parse_gemma4_args,
|
||||
Gemma4ToolParser, ToolCallDelta, ToolParseResult, ToolParser, gemma4_args,
|
||||
gemma4_array_content,
|
||||
};
|
||||
use crate::{Tool, ToolParserTestExt as _};
|
||||
use crate::Tool;
|
||||
|
||||
fn parse_gemma4_args(args: &str) -> super::Result<serde_json::Map<String, Value>> {
|
||||
let mut input = Partial::new(args);
|
||||
let _ = input.complete();
|
||||
match terminated(gemma4_args, eof).parse_next(&mut input) {
|
||||
Ok(value) => Ok(value),
|
||||
Err(ErrMode::Incomplete(_)) => Err(parsing_failed!("incomplete Gemma4 arguments")),
|
||||
Err(ErrMode::Backtrack(error) | ErrMode::Cut(error)) => {
|
||||
Err(parsing_failed!("{}", error))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_gemma4_array(array: &str) -> super::Result<Vec<Value>> {
|
||||
let mut input = array;
|
||||
let mut input = Partial::new(array);
|
||||
let _ = input.complete();
|
||||
match terminated(gemma4_array_content, eof).parse_next(&mut input) {
|
||||
Ok(value) => Ok(value),
|
||||
Err(ErrMode::Incomplete(_)) => Err(parsing_failed!("incomplete Gemma4 array")),
|
||||
@@ -494,18 +367,18 @@ mod tests {
|
||||
]
|
||||
}
|
||||
|
||||
fn collect_stream(chunks: &[&str]) -> ToolParserOutput {
|
||||
fn collect_stream(chunks: &[&str]) -> ToolParseResult {
|
||||
let mut parser = Gemma4ToolParser::new(&test_tools());
|
||||
let mut output = ToolParserOutput::default();
|
||||
let mut result = ToolParseResult::default();
|
||||
for chunk in chunks {
|
||||
output.append(parser.parse_chunk(chunk).unwrap());
|
||||
result.append(parser.push(chunk).unwrap());
|
||||
}
|
||||
output.append(parser.finish().unwrap());
|
||||
output.coalesce_calls()
|
||||
result.append(parser.finish().unwrap());
|
||||
result.coalesce_calls()
|
||||
}
|
||||
|
||||
fn first_call(output: &ToolParserOutput) -> &ToolCallDelta {
|
||||
output.calls.first().expect("expected one tool call")
|
||||
fn first_call(result: &ToolParseResult) -> &ToolCallDelta {
|
||||
result.calls.first().expect("expected one tool call")
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -543,15 +416,15 @@ mod tests {
|
||||
#[test]
|
||||
fn gemma4_parse_complete_extracts_single_tool_call() {
|
||||
let mut parser = Gemma4ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete("<|tool_call>call:get_weather{location:<|\"|>London<|\"|>}<tool_call|>")
|
||||
.unwrap();
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
|
||||
assert!(result.normal_text.is_empty());
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
|
||||
json!({ "location": "London" })
|
||||
);
|
||||
}
|
||||
@@ -568,7 +441,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn gemma4_streaming_basic_single_tool_call() {
|
||||
let output = collect_stream(&[
|
||||
let result = collect_stream(&[
|
||||
"<|tool_call>",
|
||||
"call:get_weather{",
|
||||
"location:<|\"|>Paris",
|
||||
@@ -577,17 +450,17 @@ mod tests {
|
||||
"<tool_call|>",
|
||||
]);
|
||||
|
||||
assert!(output.normal_text.is_empty());
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
|
||||
assert!(result.normal_text.is_empty());
|
||||
assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
|
||||
json!({ "location": "Paris, France" })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_streaming_text_before_and_after_tool_call() {
|
||||
let output = collect_stream(&[
|
||||
let result = collect_stream(&[
|
||||
"Let me check ",
|
||||
"the weather. ",
|
||||
"<|tool_call>",
|
||||
@@ -597,10 +470,10 @@ mod tests {
|
||||
"div>",
|
||||
]);
|
||||
|
||||
assert_eq!(output.normal_text, "Let me check the weather. <div>");
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.normal_text, "Let me check the weather. <div>");
|
||||
assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
|
||||
json!({ "location": "London" })
|
||||
);
|
||||
}
|
||||
@@ -608,68 +481,68 @@ mod tests {
|
||||
#[test]
|
||||
fn gemma4_streaming_waits_for_complete_tool_call() {
|
||||
let mut parser = Gemma4ToolParser::new(&test_tools());
|
||||
let mut output = ToolParserOutput::default();
|
||||
let mut result = ToolParseResult::default();
|
||||
|
||||
for chunk in [
|
||||
"<|tool_call>",
|
||||
"call:get_weather{",
|
||||
"location:<|\"|>Paris<|\"|>}",
|
||||
] {
|
||||
output.append(parser.parse_chunk(chunk).unwrap());
|
||||
assert!(output.calls.is_empty());
|
||||
result.append(parser.push(chunk).unwrap());
|
||||
assert!(result.calls.is_empty());
|
||||
}
|
||||
|
||||
output.append(parser.parse_chunk("<tool_call|>").unwrap());
|
||||
let output = output.coalesce_calls();
|
||||
result.append(parser.push("<tool_call|>").unwrap());
|
||||
let result = result.coalesce_calls();
|
||||
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
|
||||
json!({ "location": "Paris" })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_streaming_handles_boolean_split_across_chunks() {
|
||||
let output = collect_stream(&[
|
||||
let result = collect_stream(&[
|
||||
"<|tool_call>",
|
||||
"call:search{input:{all:tru",
|
||||
"e}}",
|
||||
"<tool_call|>",
|
||||
]);
|
||||
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("search"));
|
||||
assert_eq!(first_call(&result).name.as_deref(), Some("search"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
|
||||
json!({ "input": { "all": true } })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_streaming_handles_false_split_across_chunks() {
|
||||
let output = collect_stream(&["<|tool_call>", "call:set{flag:fals", "e}", "<tool_call|>"]);
|
||||
let result = collect_stream(&["<|tool_call>", "call:set{flag:fals", "e}", "<tool_call|>"]);
|
||||
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("set"));
|
||||
assert_eq!(first_call(&result).name.as_deref(), Some("set"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
|
||||
json!({ "flag": false })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_streaming_handles_number_split_across_chunks() {
|
||||
let output = collect_stream(&["<|tool_call>", "call:set{count:4", "2}", "<tool_call|>"]);
|
||||
let result = collect_stream(&["<|tool_call>", "call:set{count:4", "2}", "<tool_call|>"]);
|
||||
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("set"));
|
||||
assert_eq!(first_call(&result).name.as_deref(), Some("set"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
|
||||
json!({ "count": 42 })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_streaming_handles_split_string_delimiter() {
|
||||
let output = collect_stream(&[
|
||||
let result = collect_stream(&[
|
||||
"<|tool_call>",
|
||||
"call:todowrite{",
|
||||
"content:<|\"|>Buy milk<|",
|
||||
@@ -677,32 +550,17 @@ mod tests {
|
||||
"<tool_call|>",
|
||||
]);
|
||||
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("todowrite"));
|
||||
assert_eq!(first_call(&result).name.as_deref(), Some("todowrite"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
|
||||
json!({ "content": "Buy milk" })
|
||||
);
|
||||
assert!(!first_call(&output).arguments.contains("<|"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_streaming_handles_split_tool_call_end_marker() {
|
||||
let output = collect_stream(&[
|
||||
"<|tool_call>",
|
||||
"call:get_weather{location:<|\"|>Paris<|\"|>}<tool",
|
||||
"_call|>",
|
||||
]);
|
||||
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
json!({ "location": "Paris" })
|
||||
);
|
||||
assert!(!first_call(&result).arguments.contains("<|"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_streaming_handles_end_marker_literal_inside_string() {
|
||||
let output = collect_stream(&[
|
||||
let result = collect_stream(&[
|
||||
"<|tool_call>",
|
||||
"call:todowrite{",
|
||||
"content:<|\"|>literal }<tool_call|> inside",
|
||||
@@ -710,16 +568,16 @@ mod tests {
|
||||
"<tool_call|>",
|
||||
]);
|
||||
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("todowrite"));
|
||||
assert_eq!(first_call(&result).name.as_deref(), Some("todowrite"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
|
||||
json!({ "content": "literal }<tool_call|> inside" })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_streaming_handles_html_argument_without_duplication() {
|
||||
let output = collect_stream(&[
|
||||
let result = collect_stream(&[
|
||||
"<|tool_call>",
|
||||
"call:write_file{",
|
||||
"path:<|\"|>index.html<|\"|>,",
|
||||
@@ -732,9 +590,9 @@ mod tests {
|
||||
"<tool_call|>",
|
||||
]);
|
||||
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("write_file"));
|
||||
assert_eq!(first_call(&result).name.as_deref(), Some("write_file"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
|
||||
json!({
|
||||
"path": "index.html",
|
||||
"content": "<!DOCTYPE html>\n<html lang=\"zh-CN\">\n<head>\n <meta charset=\"UTF-8\">\n <meta name=\"viewport\" content=\"width=device-width\">\n",
|
||||
@@ -744,7 +602,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn gemma4_streaming_trailing_bare_bool_is_not_duplicated() {
|
||||
let output = collect_stream(&[
|
||||
let result = collect_stream(&[
|
||||
"<|tool_call>",
|
||||
"call:Edit{",
|
||||
"file_path:<|\"|>src/env.py<|\"|>,",
|
||||
@@ -755,9 +613,9 @@ mod tests {
|
||||
"<tool_call|>",
|
||||
]);
|
||||
|
||||
assert_eq!(first_call(&output).name.as_deref(), Some("Edit"));
|
||||
assert_eq!(first_call(&result).name.as_deref(), Some("Edit"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
|
||||
json!({
|
||||
"file_path": "src/env.py",
|
||||
"old_string": "old_val",
|
||||
@@ -766,7 +624,7 @@ mod tests {
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
first_call(&output).arguments.matches("replace_all").count(),
|
||||
first_call(&result).arguments.matches("replace_all").count(),
|
||||
1
|
||||
);
|
||||
}
|
||||
@@ -774,53 +632,22 @@ mod tests {
|
||||
#[test]
|
||||
fn gemma4_finish_flushes_partial_start_marker_as_text() {
|
||||
let mut parser = Gemma4ToolParser::new(&test_tools());
|
||||
let mut output = parser.parse_chunk("<").unwrap();
|
||||
output.append(parser.finish().unwrap());
|
||||
let mut result = parser.push("<").unwrap();
|
||||
result.append(parser.finish().unwrap());
|
||||
|
||||
assert_eq!(output.normal_text, "<");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(result.normal_text, "<");
|
||||
assert!(result.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_finish_rejects_complete_args_without_end_marker() {
|
||||
let mut parser = Gemma4ToolParser::new(&test_tools());
|
||||
for chunk in ["<|tool_call>", "call:get_status{}"] {
|
||||
parser.parse_chunk(chunk).unwrap();
|
||||
parser.push(chunk).unwrap();
|
||||
}
|
||||
|
||||
let error = parser.finish().unwrap_err();
|
||||
|
||||
assert!(error.to_report_string().contains("incomplete Gemma4 tool call"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_reset_preserves_internally_buffered_arguments() {
|
||||
let mut parser = Gemma4ToolParser::new(&test_tools());
|
||||
for chunk in [
|
||||
"<|tool_call>",
|
||||
"call:write_file{",
|
||||
"content:<|\"|>hello ",
|
||||
"world<|\"|>",
|
||||
] {
|
||||
parser.parse_chunk(chunk).unwrap();
|
||||
}
|
||||
|
||||
let raw = parser.reset();
|
||||
|
||||
assert_eq!(
|
||||
raw,
|
||||
"<|tool_call>call:write_file{content:<|\"|>hello world<|\"|>"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_reset_preserves_completed_arguments_after_parse_error() {
|
||||
let mut parser = Gemma4ToolParser::new(&test_tools());
|
||||
let input = "<|tool_call>call:set{broken}<tool_call|>";
|
||||
|
||||
let _error = parser.parse_chunk(input).unwrap_err();
|
||||
let raw = parser.reset();
|
||||
|
||||
assert_eq!(raw, input);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use super::{GlmXmlToolParser, Separator};
|
||||
use crate::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::{Result, Tool, ToolParseResult, ToolParser};
|
||||
|
||||
/// Tool parser for GLM-4.5/4.6 MoE XML-style tool calls.
|
||||
///
|
||||
@@ -23,6 +23,7 @@ impl Glm45MoeToolParser {
|
||||
}
|
||||
|
||||
impl ToolParser for Glm45MoeToolParser {
|
||||
/// Create a boxed GLM-4.5/4.6 MoE tool parser.
|
||||
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
@@ -30,15 +31,13 @@ impl ToolParser for Glm45MoeToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
self.0.parse_into(chunk, output)
|
||||
/// Push one decoded text chunk through the GLM MoE parser.
|
||||
fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
|
||||
self.0.push(chunk)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
/// Flush buffered text and reset parser state.
|
||||
fn finish(&mut self) -> Result<ToolParseResult> {
|
||||
self.0.finish()
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
self.0.reset()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use super::{GlmXmlToolParser, Separator};
|
||||
use crate::{Result, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::{Result, Tool, ToolParseResult, ToolParser};
|
||||
|
||||
/// Tool parser for GLM-4.7 MoE XML-style tool calls.
|
||||
///
|
||||
@@ -22,25 +22,20 @@ impl ToolParser for Glm47MoeToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
self.0.parse_into(chunk, output)
|
||||
fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
|
||||
self.0.push(chunk)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
fn finish(&mut self) -> Result<ToolParseResult> {
|
||||
self.0.finish()
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
self.0.reset()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::Glm47MoeToolParser;
|
||||
use crate::ToolParserTestExt as _;
|
||||
use super::{Glm47MoeToolParser, ToolParser};
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
|
||||
fn glm47_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
@@ -63,13 +58,13 @@ mod tests {
|
||||
)
|
||||
);
|
||||
|
||||
let output = parser.parse_complete(&output).unwrap();
|
||||
let result = parser.parse_complete(&output).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me search for that.\n");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.normal_text, "Let me search for that.\n");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({"city": "Beijing", "date": "2024-12-25"})
|
||||
);
|
||||
}
|
||||
@@ -84,14 +79,14 @@ mod tests {
|
||||
);
|
||||
|
||||
let chunks = split_by_chars(&output, 7);
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[1].name.as_deref(), Some("add"));
|
||||
assert_eq!(result.normal_text, "");
|
||||
assert_eq!(result.calls.len(), 2);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[1].name.as_deref(), Some("add"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
|
||||
json!({"x": 1, "y": 2})
|
||||
);
|
||||
}
|
||||
@@ -99,7 +94,7 @@ mod tests {
|
||||
#[test]
|
||||
fn glm47_parse_complete_converts_schema_types() {
|
||||
let mut parser = Glm47MoeToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete(&glm47_tool_call(
|
||||
"convert",
|
||||
&[
|
||||
@@ -113,7 +108,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({
|
||||
"whole": 42,
|
||||
"flag": true,
|
||||
@@ -128,12 +123,12 @@ mod tests {
|
||||
fn glm47_parse_complete_extracts_zero_argument_call() {
|
||||
let mut parser = Glm47MoeToolParser::new(&test_tools());
|
||||
|
||||
let output = parser.parse_complete("<tool_call>add</tool_call>").unwrap();
|
||||
let result = parser.parse_complete("<tool_call>add</tool_call>").unwrap();
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("add"));
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("add"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({})
|
||||
);
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@ use winnow::token::{literal, rest, take_until, take_while};
|
||||
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{parse_buffered_event, safe_text_len, xml_unescape};
|
||||
use super::{Result, ToolCallDelta, ToolParserOutput};
|
||||
use super::{Result, ToolCallDelta, ToolParseResult};
|
||||
use crate::Tool;
|
||||
|
||||
mod glm45_moe;
|
||||
@@ -76,10 +76,10 @@ impl GlmXmlToolParser {
|
||||
}
|
||||
|
||||
/// Apply one parsed GLM event to parser state and output.
|
||||
fn apply_event(&mut self, event: GlmEvent, output: &mut ToolParserOutput) -> Result<()> {
|
||||
fn apply_event(&mut self, event: GlmEvent, result: &mut ToolParseResult) -> Result<()> {
|
||||
match event {
|
||||
GlmEvent::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
result.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
}
|
||||
GlmEvent::ToolCallStart => self.mode = GlmMode::ToolCall,
|
||||
GlmEvent::ToolCall { name, raw_params } => {
|
||||
@@ -88,7 +88,7 @@ impl GlmXmlToolParser {
|
||||
let arguments = serde_json::to_string(&arguments)
|
||||
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
|
||||
|
||||
output.calls.push(ToolCallDelta {
|
||||
result.calls.push(ToolCallDelta {
|
||||
tool_index: self.emitted_tool_count,
|
||||
name: Some(name),
|
||||
arguments,
|
||||
@@ -100,36 +100,40 @@ impl GlmXmlToolParser {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
/// Reset all streaming state.
|
||||
fn reset(&mut self) {
|
||||
self.buffer.clear();
|
||||
self.mode = GlmMode::Text;
|
||||
self.emitted_tool_count = 0;
|
||||
std::mem::take(&mut self.buffer)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
/// Push one decoded text chunk through the GLM MoE parser.
|
||||
fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
|
||||
self.buffer.push_str(chunk);
|
||||
let mut result = ToolParseResult::default();
|
||||
|
||||
while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
|
||||
parse_next_glm_event(input, self.mode, self.separator)
|
||||
})? {
|
||||
self.apply_event(event, output)?;
|
||||
self.apply_event(event, &mut result)?;
|
||||
self.buffer.drain(..consumed_len);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
let mut output = ToolParserOutput::default();
|
||||
/// Flush buffered text and reset parser state.
|
||||
fn finish(&mut self) -> Result<ToolParseResult> {
|
||||
let mut result = ToolParseResult::default();
|
||||
if !self.buffer.is_empty() {
|
||||
match self.mode {
|
||||
GlmMode::Text => output.normal_text.push_str(&self.buffer),
|
||||
GlmMode::Text => result.normal_text.push_str(&self.buffer),
|
||||
GlmMode::ToolCall => return Err(parsing_failed!("incomplete GLM MoE tool call")),
|
||||
GlmMode::AfterToolCall => {}
|
||||
}
|
||||
}
|
||||
let _ = self.reset();
|
||||
Ok(output)
|
||||
self.reset();
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -252,8 +256,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::Glm45MoeToolParser;
|
||||
use crate::ToolParser;
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParser, ToolParserTestExt as _};
|
||||
|
||||
fn glm45_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
let params = params
|
||||
@@ -269,10 +273,10 @@ mod tests {
|
||||
#[test]
|
||||
fn glm45_parse_complete_without_tool_call_keeps_text() {
|
||||
let mut parser = Glm45MoeToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("Hello, world!").unwrap();
|
||||
let result = parser.parse_complete("Hello, world!").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Hello, world!");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(result.normal_text, "Hello, world!");
|
||||
assert!(result.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -286,13 +290,13 @@ mod tests {
|
||||
)
|
||||
);
|
||||
|
||||
let output = parser.parse_complete(&output).unwrap();
|
||||
let result = parser.parse_complete(&output).unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Let me search for that.\n");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.normal_text, "Let me search for that.\n");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({"city": "Beijing", "date": "2024-12-25"})
|
||||
);
|
||||
}
|
||||
@@ -307,14 +311,14 @@ mod tests {
|
||||
);
|
||||
|
||||
let chunks = split_by_chars(&output, 11);
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
let result = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(output.calls[1].name.as_deref(), Some("add"));
|
||||
assert_eq!(result.normal_text, "");
|
||||
assert_eq!(result.calls.len(), 2);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.calls[1].name.as_deref(), Some("add"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
|
||||
json!({"x": 1, "y": 2})
|
||||
);
|
||||
}
|
||||
@@ -322,7 +326,7 @@ mod tests {
|
||||
#[test]
|
||||
fn glm45_parse_complete_unescapes_literal_closing_tags_in_arg_value() {
|
||||
let mut parser = Glm45MoeToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
let result = parser
|
||||
.parse_complete(&glm45_tool_call(
|
||||
"get_weather",
|
||||
&[
|
||||
@@ -333,7 +337,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
|
||||
serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
|
||||
json!({
|
||||
"city": "Paris </arg_value></tool_call>",
|
||||
"date": "2026-05-08",
|
||||
@@ -345,17 +349,17 @@ mod tests {
|
||||
fn glm45_streaming_without_tool_call_emits_text_incrementally() {
|
||||
let mut parser = Glm45MoeToolParser::new(&test_tools());
|
||||
|
||||
let output = collect_stream(&mut parser, &["hello ", "world"]);
|
||||
let result = collect_stream(&mut parser, &["hello ", "world"]);
|
||||
|
||||
assert_eq!(output.normal_text, "hello world");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(result.normal_text, "hello world");
|
||||
assert!(result.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn glm45_streaming_preserves_prefix_text() {
|
||||
let mut parser = Glm45MoeToolParser::new(&test_tools());
|
||||
|
||||
let output = collect_stream(
|
||||
let result = collect_stream(
|
||||
&mut parser,
|
||||
&[
|
||||
"Prefix ",
|
||||
@@ -363,14 +367,14 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "Prefix ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(result.normal_text, "Prefix ");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn glm45_streaming_handles_start_token_split_across_chunks() {
|
||||
let mut parser = Glm45MoeToolParser::new(&test_tools());
|
||||
let output = collect_stream(
|
||||
let result = collect_stream(
|
||||
&mut parser,
|
||||
&[
|
||||
"hello <tool",
|
||||
@@ -379,26 +383,26 @@ mod tests {
|
||||
],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(result.normal_text, "hello ");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn glm45_streaming_does_not_emit_incomplete_tool_call() {
|
||||
let mut parser = Glm45MoeToolParser::new(&test_tools());
|
||||
|
||||
let output = parser.parse_chunk("<tool_call>get_weather\n<arg_key>city</arg_key>").unwrap();
|
||||
let result = parser.push("<tool_call>get_weather\n<arg_key>city</arg_key>").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert!(output.calls.is_empty());
|
||||
assert_eq!(result.normal_text, "");
|
||||
assert!(result.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn glm45_finish_fails_incomplete_tool_call() {
|
||||
let mut parser = Glm45MoeToolParser::new(&test_tools());
|
||||
|
||||
parser.parse_chunk("<tool_call>get_weather\n<arg_key>city</arg_key>").unwrap();
|
||||
parser.push("<tool_call>get_weather\n<arg_key>city</arg_key>").unwrap();
|
||||
let error = parser.finish().unwrap_err();
|
||||
|
||||
assert!(error.as_report().to_string().contains("incomplete GLM MoE tool call"));
|
||||
@@ -408,7 +412,7 @@ mod tests {
|
||||
fn glm45_malformed_tool_call_fails_fast() {
|
||||
let mut parser = Glm45MoeToolParser::new(&test_tools());
|
||||
|
||||
let error = parser.parse_chunk("<tool_call>get_weather<arg_key>city</arg_key><arg_value>Paris</arg_value></tool_call>").unwrap_err();
|
||||
let error = parser.push("<tool_call>get_weather<arg_key>city</arg_key><arg_value>Paris</arg_value></tool_call>").unwrap_err();
|
||||
|
||||
assert!(error.as_report().to_string().contains("tool parser parsing failed"));
|
||||
}
|
||||
@@ -417,7 +421,7 @@ mod tests {
|
||||
fn glm45_streaming_ignores_trailing_text_after_tool_calls() {
|
||||
let mut parser = Glm45MoeToolParser::new(&test_tools());
|
||||
|
||||
let output = collect_stream(
|
||||
let result = collect_stream(
|
||||
&mut parser,
|
||||
&[&format!(
|
||||
"{}<|endoftext|>",
|
||||
@@ -425,7 +429,7 @@ mod tests {
|
||||
)],
|
||||
);
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(result.normal_text, "");
|
||||
assert_eq!(result.calls.len(), 1);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,559 +0,0 @@
|
||||
use winnow::ascii::multispace0 as ws0;
|
||||
use winnow::combinator::{alt, delimited, eof, repeat, seq, terminated};
|
||||
use winnow::prelude::*;
|
||||
use winnow::stream::Partial;
|
||||
use winnow::token::{literal, rest, take_until};
|
||||
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{parse_buffered_event, safe_text_len};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::Tool;
|
||||
|
||||
const TOOL_CALLS_START: &str = "<tool_calls>";
|
||||
const TOOL_CALLS_END: &str = "</tool_calls>";
|
||||
const TOOL_CALL_START: &str = "<tool_call>";
|
||||
const TOOL_CALL_END: &str = "</tool_call>";
|
||||
const TOOL_SEP: &str = "<tool_sep>";
|
||||
const ARG_KEY_START: &str = "<arg_key>";
|
||||
const ARG_KEY_END: &str = "</arg_key>";
|
||||
const ARG_VALUE_START: &str = "<arg_value>";
|
||||
const ARG_VALUE_END: &str = "</arg_value>";
|
||||
|
||||
type HyV3Input<'i> = Partial<&'i str>;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum HyV3Mode {
|
||||
Text,
|
||||
ToolBlock,
|
||||
Done,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum HyV3Event {
|
||||
Text {
|
||||
len: usize,
|
||||
},
|
||||
ToolBlockStart,
|
||||
ToolCall {
|
||||
name: String,
|
||||
raw_params: Vec<(String, String)>,
|
||||
},
|
||||
ToolBlockEnd,
|
||||
IgnoredRest,
|
||||
}
|
||||
|
||||
/// Tool parser for HY3 XML-style tool calls.
|
||||
///
|
||||
/// Example tool call content:
|
||||
///
|
||||
/// ```text
|
||||
/// <tool_calls>
|
||||
/// <tool_call>get_weather<tool_sep>
|
||||
/// <arg_key>city</arg_key><arg_value>Beijing</arg_value>
|
||||
/// </tool_call>
|
||||
/// </tool_calls>
|
||||
/// ```
|
||||
///
|
||||
/// Arguments are emitted only after a full `<tool_call>` block is parsed.
|
||||
/// HY3 marker tokens are added-vocabulary tokens rather than tokenizer special
|
||||
/// tokens, so the default `preserve_special_tokens() == false` is sufficient.
|
||||
pub struct HyV3ToolParser {
|
||||
buffer: String,
|
||||
mode: HyV3Mode,
|
||||
emitted_tool_count: usize,
|
||||
tool_parameters: ToolSchemas,
|
||||
}
|
||||
|
||||
impl HyV3ToolParser {
|
||||
/// Create a HY3 tool parser.
|
||||
fn new(tools: &[Tool]) -> Self {
|
||||
Self {
|
||||
buffer: String::new(),
|
||||
mode: HyV3Mode::Text,
|
||||
emitted_tool_count: 0,
|
||||
tool_parameters: ToolSchemas::from_tools(tools),
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply one parsed HY3 event to parser state and output.
|
||||
fn apply_event(&mut self, event: HyV3Event, output: &mut ToolParserOutput) -> Result<()> {
|
||||
match event {
|
||||
HyV3Event::Text { len: consumed_len } => {
|
||||
output.normal_text.push_str(&self.buffer[..consumed_len]);
|
||||
}
|
||||
HyV3Event::ToolBlockStart => self.mode = HyV3Mode::ToolBlock,
|
||||
HyV3Event::ToolCall { name, raw_params } => {
|
||||
let arguments = self.tool_parameters.convert_params_with_schema(&name, raw_params);
|
||||
let arguments = serde_json::to_string(&arguments)
|
||||
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
|
||||
|
||||
output.calls.push(ToolCallDelta {
|
||||
tool_index: self.emitted_tool_count,
|
||||
name: Some(name),
|
||||
arguments,
|
||||
});
|
||||
self.emitted_tool_count += 1;
|
||||
}
|
||||
HyV3Event::ToolBlockEnd => self.mode = HyV3Mode::Done,
|
||||
HyV3Event::IgnoredRest => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolParser for HyV3ToolParser {
|
||||
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
self.buffer.push_str(chunk);
|
||||
|
||||
while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
|
||||
parse_next_hy_v3_event(input, self.mode)
|
||||
})? {
|
||||
self.apply_event(event, output)?;
|
||||
self.buffer.drain(..consumed_len);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
let mut output = ToolParserOutput::default();
|
||||
match self.mode {
|
||||
HyV3Mode::Text => output.normal_text.push_str(&self.buffer),
|
||||
HyV3Mode::ToolBlock => return Err(parsing_failed!("incomplete HY3 tool call")),
|
||||
HyV3Mode::Done => {}
|
||||
}
|
||||
let _ = self.reset();
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
self.mode = HyV3Mode::Text;
|
||||
self.emitted_tool_count = 0;
|
||||
std::mem::take(&mut self.buffer)
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a HY3 event for the current parser mode.
|
||||
fn parse_next_hy_v3_event(input: &mut HyV3Input<'_>, mode: HyV3Mode) -> ModalResult<HyV3Event> {
|
||||
match mode {
|
||||
HyV3Mode::Text => parse_text_event(input),
|
||||
HyV3Mode::ToolBlock => parse_tool_block_event(input),
|
||||
HyV3Mode::Done => ignored_rest_event(input),
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a text-mode HY3 event.
|
||||
fn parse_text_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
|
||||
alt((tool_block_start_event, safe_text_event)).parse_next(input)
|
||||
}
|
||||
|
||||
/// Parse a HY3 tool-block start marker.
|
||||
fn tool_block_start_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
|
||||
literal(TOOL_CALLS_START).value(HyV3Event::ToolBlockStart).parse_next(input)
|
||||
}
|
||||
|
||||
/// Parse a safe text run before the next HY3 marker.
|
||||
fn safe_text_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
|
||||
safe_text_len(input, TOOL_CALLS_START).map(|len| HyV3Event::Text { len })
|
||||
}
|
||||
|
||||
/// Parse one event inside a HY3 tool block.
|
||||
fn parse_tool_block_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
|
||||
alt((tool_block_end_event, tool_call_event)).parse_next(input)
|
||||
}
|
||||
|
||||
/// Parse a HY3 tool-block end marker.
|
||||
fn tool_block_end_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
|
||||
(ws0, literal(TOOL_CALLS_END)).value(HyV3Event::ToolBlockEnd).parse_next(input)
|
||||
}
|
||||
|
||||
/// Parse a complete HY3 tool-call block.
|
||||
fn tool_call_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
|
||||
let (name, body) = seq!(
|
||||
_: ws0,
|
||||
_: literal(TOOL_CALL_START),
|
||||
take_until(0.., TOOL_SEP),
|
||||
_: literal(TOOL_SEP),
|
||||
take_until(0.., TOOL_CALL_END),
|
||||
_: literal(TOOL_CALL_END),
|
||||
)
|
||||
.parse_next(input)?;
|
||||
let raw_params = parse_tool_call_params(body)?;
|
||||
|
||||
Ok(HyV3Event::ToolCall {
|
||||
name: name.trim().to_string(),
|
||||
raw_params,
|
||||
})
|
||||
}
|
||||
|
||||
/// Parse all parameter blocks inside a complete HY3 tool call.
|
||||
fn parse_tool_call_params(tool_call_body: &str) -> ModalResult<Vec<(String, String)>> {
|
||||
let mut input = tool_call_body;
|
||||
delimited(ws0, repeat(0.., terminated(parameter, ws0)), eof).parse_next(&mut input)
|
||||
}
|
||||
|
||||
/// Parse a HY3 argument key/value block.
|
||||
fn parameter(input: &mut &str) -> ModalResult<(String, String)> {
|
||||
let (name, value) = seq!(
|
||||
_: literal(ARG_KEY_START),
|
||||
take_until(0.., ARG_KEY_END),
|
||||
_: literal(ARG_KEY_END),
|
||||
_: ws0,
|
||||
_: literal(ARG_VALUE_START),
|
||||
take_until(0.., ARG_VALUE_END),
|
||||
_: literal(ARG_VALUE_END),
|
||||
)
|
||||
.parse_next(input)?;
|
||||
|
||||
Ok((name.trim().to_string(), value.to_string()))
|
||||
}
|
||||
|
||||
/// Parse ignored rest after the HY3 tool block ends.
|
||||
fn ignored_rest_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
|
||||
rest.value(HyV3Event::IgnoredRest).parse_next(input)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use expect_test::expect;
|
||||
use serde_json::{Value, json};
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::{HyV3ToolParser, ToolParser};
|
||||
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::{ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
let params = params
|
||||
.iter()
|
||||
.map(|(name, value)| format!("<arg_key>{name}</arg_key><arg_value>{value}</arg_value>"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
format!("<tool_call>{function_name}<tool_sep>{params}</tool_call>")
|
||||
}
|
||||
|
||||
fn build_tool_calls(tool_calls: &[String]) -> String {
|
||||
format!("<tool_calls>\n{}\n</tool_calls>", tool_calls.join("\n"))
|
||||
}
|
||||
|
||||
fn parsed_arguments(output: &ToolParserOutput, index: usize) -> Value {
|
||||
serde_json::from_str(&output.calls[index].arguments).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_does_not_preserve_special_tokens() {
|
||||
let parser = HyV3ToolParser::new(&test_tools());
|
||||
|
||||
assert!(!parser.preserve_special_tokens());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_parse_complete_without_tool_call_keeps_text() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let output = parser.parse_complete("This is a plain response.").unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "This is a plain response.");
|
||||
assert!(output.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_parse_complete_extracts_zero_arg_inline_tool_call() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
.parse_complete(
|
||||
"<tool_calls><tool_call>get_current_date<tool_sep></tool_call></tool_calls>",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
|
||||
assert_eq!(parsed_arguments(&output, 0), json!({}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_parse_complete_extracts_zero_arg_newline_tool_call() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
.parse_complete(
|
||||
"<tool_calls>\n<tool_call>get_current_date<tool_sep>\n</tool_call>\n</tool_calls>",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
|
||||
assert_eq!(parsed_arguments(&output, 0), json!({}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_parse_complete_extracts_arguments_on_same_line() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
.parse_complete(
|
||||
"<tool_calls><tool_call>get_weather<tool_sep><arg_key>city</arg_key><arg_value>Beijing</arg_value><arg_key>date</arg_key><arg_value>2026-03-30</arg_value></tool_call></tool_calls>",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
parsed_arguments(&output, 0),
|
||||
json!({ "city": "Beijing", "date": "2026-03-30" })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_parse_complete_extracts_arguments_with_newlines() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
.parse_complete(&build_tool_calls(&[build_tool_call(
|
||||
"get_weather",
|
||||
&[("city", "Beijing"), ("date", "2026-03-30")],
|
||||
)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
parsed_arguments(&output, 0),
|
||||
json!({ "city": "Beijing", "date": "2026-03-30" })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_parse_complete_preserves_prefix_and_ignores_trailing_text() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
.parse_complete(&format!(
|
||||
"Checking.{} trailing text",
|
||||
build_tool_calls(&[build_tool_call("get_current_date", &[])])
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "Checking.");
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_parse_complete_extracts_multiple_tool_calls_in_one_block() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
.parse_complete(&build_tool_calls(&[
|
||||
build_tool_call(
|
||||
"get_weather",
|
||||
&[("city", "Beijing"), ("date", "2026-03-30")],
|
||||
),
|
||||
build_tool_call(
|
||||
"get_weather",
|
||||
&[("city", "Hangzhou"), ("date", "2026-03-30")],
|
||||
),
|
||||
]))
|
||||
.unwrap();
|
||||
|
||||
expect![[r#"
|
||||
ToolParserOutput {
|
||||
normal_text: "",
|
||||
calls: [
|
||||
ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"city\":\"Beijing\",\"date\":\"2026-03-30\"}",
|
||||
},
|
||||
ToolCallDelta {
|
||||
tool_index: 1,
|
||||
name: Some(
|
||||
"get_weather",
|
||||
),
|
||||
arguments: "{\"city\":\"Hangzhou\",\"date\":\"2026-03-30\"}",
|
||||
},
|
||||
],
|
||||
}
|
||||
"#]]
|
||||
.assert_debug_eq(&output);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_parse_complete_converts_schema_types() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let output = parser
|
||||
.parse_complete(&build_tool_calls(&[build_tool_call(
|
||||
"convert",
|
||||
&[
|
||||
("whole", "5.3"),
|
||||
("flag", "true"),
|
||||
("payload", r#"{"k":1}"#),
|
||||
("items", "[1,2]"),
|
||||
],
|
||||
)]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
parsed_arguments(&output, 0),
|
||||
json!({
|
||||
"whole": 5.3,
|
||||
"flag": true,
|
||||
"payload": { "k": 1 },
|
||||
"items": [1, 2]
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_streaming_without_tool_call_emits_text_incrementally() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let mut output = ToolParserOutput::default();
|
||||
|
||||
output.append(parser.parse_chunk("This is ").unwrap());
|
||||
output.append(parser.parse_chunk("a plain ").unwrap());
|
||||
output.append(parser.parse_chunk("response.").unwrap());
|
||||
output.append(parser.finish().unwrap());
|
||||
|
||||
assert_eq!(output.normal_text, "This is a plain response.");
|
||||
assert!(output.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_streaming_extracts_zero_arg_tool_call() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let chunks = [
|
||||
"<tool_calls>",
|
||||
"\n<tool_call>",
|
||||
"get_current_date",
|
||||
"<tool_sep>",
|
||||
"\n</tool_call>",
|
||||
"\n</tool_calls>",
|
||||
];
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
|
||||
assert_eq!(parsed_arguments(&output, 0), json!({}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_streaming_extracts_arguments() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let chunks = [
|
||||
"<tool_calls>",
|
||||
"\n<tool_call>",
|
||||
"get_weather",
|
||||
"<tool_sep>",
|
||||
"\n<arg_key>city</arg_key>",
|
||||
"\n<arg_value>Beijing</arg_value>",
|
||||
"\n<arg_key>date</arg_key>",
|
||||
"\n<arg_value>2026-03-30</arg_value>",
|
||||
"\n</tool_call>",
|
||||
"\n</tool_calls>",
|
||||
];
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
|
||||
assert_eq!(
|
||||
parsed_arguments(&output, 0),
|
||||
json!({ "city": "Beijing", "date": "2026-03-30" })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_streaming_preserves_prefix_text() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let chunks = [
|
||||
"Checking.",
|
||||
"<tool_calls>",
|
||||
"\n<tool_call>",
|
||||
"get_current_date",
|
||||
"<tool_sep>",
|
||||
"\n</tool_call>",
|
||||
"\n</tool_calls>",
|
||||
];
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "Checking.");
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_streaming_extracts_multiple_tool_calls_in_one_block() {
|
||||
let input = build_tool_calls(&[
|
||||
build_tool_call(
|
||||
"get_weather",
|
||||
&[("city", "Beijing"), ("date", "2026-03-30")],
|
||||
),
|
||||
build_tool_call(
|
||||
"get_weather",
|
||||
&[("city", "Hangzhou"), ("date", "2026-03-30")],
|
||||
),
|
||||
]);
|
||||
let chunks = split_by_chars(&input, 9);
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.calls.len(), 2);
|
||||
assert_eq!(parsed_arguments(&output, 0)["city"], json!("Beijing"));
|
||||
assert_eq!(parsed_arguments(&output, 1)["city"], json!("Hangzhou"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_streaming_handles_start_marker_split_across_chunks() {
|
||||
let input = format!(
|
||||
"hello {}",
|
||||
build_tool_calls(&[build_tool_call("get_weather", &[("city", "Beijing")])])
|
||||
);
|
||||
let chunks = split_by_chars(&input, 5);
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
|
||||
let output = collect_stream(&mut parser, &chunks);
|
||||
|
||||
assert_eq!(output.normal_text, "hello ");
|
||||
assert_eq!(output.calls.len(), 1);
|
||||
assert_eq!(parsed_arguments(&output, 0), json!({ "city": "Beijing" }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_streaming_does_not_emit_incomplete_tool_call() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let mut output = ToolParserOutput::default();
|
||||
|
||||
parser
|
||||
.parse_into(
|
||||
"<tool_calls><tool_call>get_weather<tool_sep><arg_key>city</arg_key><arg_value>Bei",
|
||||
&mut output,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output.normal_text, "");
|
||||
assert!(output.calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_finish_fails_incomplete_tool_call() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
parser.parse_chunk("<tool_calls><tool_call>get_weather<tool_sep>").unwrap();
|
||||
|
||||
let error = parser.finish().unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: incomplete HY3 tool call"]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hy_v3_malformed_tool_call_fails_fast() {
|
||||
let mut parser = HyV3ToolParser::new(&test_tools());
|
||||
let error = parser
|
||||
.parse_complete(
|
||||
"<tool_calls><tool_call>get_weather<tool_sep><arg_key>city</arg_key><arg_value>Beijing</tool_call></tool_calls>",
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(error.to_report_string().starts_with("tool parser parsing failed:"));
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user