Compare commits

..
Author SHA1 Message Date
Bugen Zhao 6e714a103c stash
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-02 05:51:48 +00:00
Bugen Zhao c9951fd5c7 separate vllm-model-files
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-01 14:24:43 +00:00
Harry MellorandGitHub a78c15616f Migrate GPTBigCode and Starcoder2 to the Transformers modeling backend (#30966)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-01 13:41:36 +00:00
5c4db60f01 docs(security): document gRPC interface as insecure for private use only (#45903)
Signed-off-by: jperezde <jperezde@redhat.com>
Signed-off-by: Russell Bryant <russell.bryant@gmail.com>
Co-authored-by: Russell Bryant <russell.bryant@gmail.com>
Co-authored-by: Russell Bryant <rbryant@redhat.com>
2026-07-01 12:39:57 +00:00
4e5ca89cfe [ROCm][MiniMax-M3] Cross-layer lightning-indexer top-k sharing (#47269)
Signed-off-by: Fangzhou Ai <fangzhouai@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-01 10:50:09 +00:00
Harry MellorandGitHub a22e0dfc69 [Model] Remove AyaVision, MusicFlamingo (#47263)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-01 10:39:33 +00:00
stevenkuangandGitHub cc56379e28 [Model] Support Hy3 token suffix and JSON Schema array types (#47192)
Signed-off-by: stevenkuang-tencent <stevenkuang@tencent.com>
2026-07-01 10:16:07 +00:00
024b06b0dc [Bugfix] Expose usage field in GenerateResponse for disaggregated serving (#42748)
Signed-off-by: AIvashov <ivashov.aleksey@proton.me>
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Co-authored-by: NickLucche <nicolo.lucchesi@mistral.ai>
2026-07-01 10:00:19 +00:00
Harry MellorandGitHub e7d0fcbc09 [CI] Fix various failures on main (#47197)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-01 10:35:34 +01:00
akii96andGitHub aa8bb5562e [ROCm][Perf][Bugfix] DSv4 indexer: use platform FP8 dtype (fnuz) for Q-quant on gfx942 (#46730)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
2026-07-01 17:33:55 +08:00
Andy LoandGitHub fa4bec9056 [Bugfix] Fix pooled Whisper sliding-window KV sizing (#47071)
Signed-off-by: Andy Lo <andy@mistral.ai>
2026-07-01 11:33:19 +02:00
dee5da1dec [Test] Run SageMaker handler-override tests in-process via TestClient (#47250)
Signed-off-by: Jyothirmai Kottu <jkottu@amazon.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-01 09:14:00 +00:00
ed41aa270a [ROCm][DSV4] Use aiter mHC pre/post as the default ROCm path (#43950)
Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
Signed-off-by: Fangzhou-Ai <fangzhouai@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-01 16:27:42 +08:00
77a9c5ae28 Weight sync refactor + move sparse nccl engine (#44353)
Signed-off-by: hao-aaron <ahao@anyscale.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-07-01 01:25:19 -07:00
f651a8a9a4 [XPU][UT]Enable ut qk_norm_rope_fusion (#42486)
Signed-off-by: Lai, Yejing <yejing.lai@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-01 07:38:03 +00:00
Jee Jee LiandGitHub 8f82be5705 [CI/Build] Fix LoRA testing (#47242)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
2026-07-01 15:36:13 +08:00
Nils MattesonandGitHub a461070d1c [Core] Make sleep-mode backend capability flags communicator-agnostic (#47243) 2026-07-01 07:17:44 +00:00
4470ae84de Remove mantis (#46806)
Signed-off-by: Xianbao QIAN <xianbao.qian@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-01 07:13:58 +00:00
ChaunceyandGitHub 697c34b97b [Bugfix] Fix beam search candidate indexing when logprobs count varies (#47126)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
2026-07-01 07:07:06 +00:00
Blas Rodriguez IrizarandGitHub 5b431b905c [Rust Frontend] Coerce completion max_tokens: null to default (#47166)
Signed-off-by: Blas Rodriguez Irizar <rodrigblas@gmail.com>
2026-07-01 06:41:33 +00:00
almayneGitHubLi, Jiang <jiang1.li@intel.com>
89e99202f2 [CPU][Perf]Added tanh AOR for faster gelu activations. (#44639)
Signed-off-by: Anna Mayne <anna.mayne@arm.com>
Signed-off-by: almayne <anna.mayne@arm.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-06-30 23:24:40 -07:00
Micah WilliamsonandGitHub b446792306 [ROCm][Bugfix] Fix Triton "out of resource: shared memory" Error In One-Shot LoRA MoE (#47209)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-06-30 23:24:36 -07:00
Micah WilliamsonandGitHub c3b1f9e827 [ROCm][CI] Enable LoRA TP Distributed Test Group In AMD CI (#47193)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-06-30 23:24:32 -07:00
Jonathan MamouandGitHub df802a87b7 [CPU] Remove speculative decoding stream overrides from CPUModelRunner (#47162)
Signed-off-by: jmamou <jonathan.mamou@intel.com>
2026-07-01 06:12:49 +00:00
Nils MattesonandGitHub 93d8f834dd [Core] Pluggable sleep-mode backend abstraction (RFC #34303) (#44074) 2026-06-30 22:00:53 -07:00
Maria GuevaraandGitHub aeb35b90f0 [Rust Frontend] Add error context in tool parser failures (#46512)
Signed-off-by: Maria Guevara <kawaiiplush14@gmail.com>
2026-07-01 12:48:55 +08:00
Gabriel WuandGitHub 9a08a5118e fix: skip cooperative top-K on SM120 (#47164)
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
2026-06-30 21:32:54 -07:00
c5200d3565 [Attention][DSA] support dcp for FLASHINFER_MLA_SPARSE (#46076)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Jingyi Yang <girasoleyang@gmail.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: GirasoleY <girasoleyang@gmail.com>
Co-authored-by: Jingyi Yang <girasoleyang@gmail.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-01 00:32:20 -04:00
MattandGitHub 3c1396bab6 [Hardware][AMD][CI] Toggle test coredumps on ROCm debug agent (#47222)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
2026-06-30 23:30:10 -05:00
Benjamin ChislettandGitHub 9969466a59 [Spec Decode] Support SWA + DFlash for MiMo (#46104) 2026-06-30 20:34:47 -07:00
achyuthan.sandGitHub 3406e8f83d [Bugfix][Frontend][gpt-oss] Return raw output when Harmony parser ends non-terminal (#47062)
Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
2026-07-01 01:46:01 +00:00
a264e41975 [Distributed] Default FlashInfer allreduce to mnnvl on single node (#47219)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-30 18:35:56 -07:00
Woosuk KwonandGitHub f098ee70c7 [GLM5] Support FlashMLA FP8 KV cache (Hopper & Blackwell) (#47090)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-06-30 18:13:21 -07:00
9294dd27eb fix(reasoning): guard rfind in ernie45 streaming </response> branch (#46255)
Signed-off-by: Chenglun Hu <chenglunhu@gmail.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>
2026-07-01 01:01:14 +00:00
yzong-rhandGitHub b1190d03cc [Refactor][GPT-OSS] Harmony Responses API Refactor to use HarmonyParser (#47185)
Signed-off-by: Yifan Zong <yzong@redhat.com>
2026-06-30 19:23:20 -04:00
92c7fac640 [Perf] Restore zero-init of swizzled NVFP4 scale buffer to recover Blackwell decode throughput (#45739)
Signed-off-by: Albert Cheng <albertching0112@gmail.com>
Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
2026-06-30 22:56:56 +00:00
Ting SUNandGitHub ac521f6237 [Bugfix][Structured Outputs] Reject degenerate structured_outputs that crash EngineCore (#45346)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
2026-06-30 22:41:33 +00:00
28242824e0 [Bugfix][Frontend] Normalize constrained Harmony recipients (#45657)
Signed-off-by: shaojunjie <626650687@qq.com>
Co-authored-by: Ben Browning <bbrownin@redhat.com>
2026-06-30 17:33:10 -04:00
VectorPeakandGitHub 68294739d1 [Bugfix] Align OpenCV video metadata timeline (#47099)
Signed-off-by: VectorPeak <73048950+VectorPeak@users.noreply.github.com>
2026-06-30 20:43:42 +00:00
c8d2f3cb14 [Bugfix] compressed-tensors: allow int8 grouped WNA16 MoE on Marlin (#47154)
Signed-off-by: Joe Rowell <joerowell4@gmail.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
2026-06-30 12:50:46 -07:00
MattandGitHub 345b28ff2f [Hardware][AMD][CI] Bump timeouts of various test groups on AMD CI (#47195)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
2026-06-30 14:30:53 -05:00
248d1fbb71 [Feat][1/N] CuTeDSL warmup infrastructure, FA4 MLA (#46182)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
2026-06-30 12:17:34 -07:00
11b26c5528 [Bugfix][Tool Parser] PoolsideV1: fix logprobs AttributeError on Responses API (#47138)
Signed-off-by: Joe Rowell <joerowell4@gmail.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
2026-06-30 19:14:09 +00:00
Roberto L. CastroandGitHub 20434c472e [Feat] Improve Triton JIT diagnostics (#46621)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
2026-06-30 18:50:15 +00:00
Andreas KaratzasandGitHub c8f9c156a5 [ROCm][V1][MLA] Clone prefill backend state per metadata builder (#46993)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-06-30 11:43:54 -07:00
953bba488d [PERF] Extend NCCL symmetric memory to AllGather and ReduceScatter (#46703)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: snordmann <snordmann@nvidia.com>
2026-06-30 11:38:18 -07:00
Wentao YeandGitHub 3a9784b82c [Feature] DP supervisor using rust frontend (#47076)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-30 14:34:05 -04:00
Giancarlo DelfinandGitHub 3cecee40f3 [Model Runner V2][Spec Decode] Fix stale values in idx_mapping from CG num reqs padding (#47066) 2026-06-30 11:25:32 -07:00
a7732537f4 [Bugfix] Restore part of bugfix #42650 after accidental deletion in #43241 (#47039)
Signed-off-by: zhanda <zhandazhu@gmail.com>
Signed-off-by: Nikita Shapovalov <nikita@poolside.ai>
Co-authored-by: Zhanda Zhu <49645678+zhandaz@users.noreply.github.com>
Co-authored-by: Shang Wang <shangw@nvidia.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
2026-06-30 11:07:59 -07:00
Rishi PuriGitHubAnshika OjhaClaudegemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>Stefano Castagnetta
727971f1c1 Add Medusa speculative decoding e2e test (#41396)
Signed-off-by: Anshika Ojha <anshikao@nvidia.com>
Signed-off-by: Rishi Puri <riship@nvidia.com>
Signed-off-by: Rishi Puri <puririshi98@berkeley.edu>
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
Co-authored-by: Anshika Ojha <215760622+ojhaanshika@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Stefano Castagnetta <scastagnetta@nvidia.com>
2026-06-30 18:02:22 +00:00
25671cb520 [Parser][Bugfix] Ensure tool call or other special tokens don't leak in non-streaming tool parsing (#46875)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-06-30 13:46:53 -04:00
Lucas WilkinsonGitHubOpenAI Codexmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
27d5f78b63 [CI] Move distributed small LM eval to B200 (#47048)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-30 13:34:25 -04:00
liuzhenweiandGitHub 7a341fa109 [XPU] Support ZE_AFFINITY_MASK passthrough in xpu_disagg_acc_test (#47105)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
2026-06-30 17:06:12 +00:00
Charlie FuandGitHub f41e8ddc97 [ROCm][CI] Move PyTorch Compilation Unit Tests to MI300(gfx942) (#47065)
Signed-off-by: charlifu <charlifu@amd.com>
2026-06-30 11:32:58 -05:00
fangyuchuGitHubTyler Michael Smithmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
245888ff77 [Feature] Detect all2all peer fault with fault tolerance backend and prevent corrupted output (#43637)
Signed-off-by: fangyuchu <fangyuchu@qq.com>
Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-30 09:00:25 -07:00
e840f0d3f5 [Platform] Replace torch.cuda.Event with torch.Event (#47140)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-06-30 08:39:59 -07:00
fcaa84efa7 [BugFix] Gate MRV2 mixed sparse-MLA warmup on max_num_seqs > 1 (#47050)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: ziminghuang <ziminghuang@inferact.ai>
2026-06-30 16:31:27 +01:00
Wentao YeandGitHub 9e84ec8648 [Refactor] Remove dead minimax allreduce rms kernel (#46842)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-30 08:29:21 -07:00
252 changed files with 8204 additions and 5821 deletions
-25
View File
@@ -1,25 +0,0 @@
group: Zen CPU
depends_on: []
steps:
- label: Zen-CPU-Kernel Tests
depends_on: []
soft_fail: false
device: zen5
no_plugin: true
source_file_dependencies:
- setup.py
- docker/Dockerfile.cpu
- requirements/cpu.txt
- requirements/build/cpu.txt
- csrc/cpu/
- cmake/cpu_extension.cmake
- CMakeLists.txt
- vllm/model_executor/layers/utils.py
- vllm/platforms/cpu.py
- vllm/platforms/zen_cpu.py
- vllm/platforms/__init__.py
- tests/model_executor/test_cpu_unquantized_gemm_dispatch.py
commands:
- |
bash .buildkite/scripts/hardware_ci/run-zen-cpu-test.sh 20m "
pytest -x -v -s tests/model_executor/test_cpu_unquantized_gemm_dispatch.py"
+25
View File
@@ -103,6 +103,31 @@ steps:
pytest -v -s v1/kv_offload &&
pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
- label: NixlConnector PD accuracy (2 GPUs)
timeout_in_minutes: 60
num_devices: 2
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 24+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/xpu.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
- label: Regression
key: regression
timeout_in_minutes: 30
@@ -22,7 +22,7 @@ steps:
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
'pip install av &&
cd tests &&
pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" &&
pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model'
@@ -47,8 +47,7 @@ steps:
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
cd tests &&
'cd tests &&
pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model'
- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
@@ -71,8 +70,7 @@ steps:
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
cd tests &&
'cd tests &&
pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" &&
pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model'
@@ -96,7 +94,7 @@ steps:
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
'pip install av &&
cd tests &&
pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing'
@@ -121,7 +119,7 @@ steps:
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install av matplotlib ftfy git+https://github.com/TIGER-AI-Lab/Mantis.git &&
'pip install av matplotlib ftfy &&
pip install open-clip-torch --no-deps &&
cd tests &&
pytest -v -s models/multimodal/processing/test_tensor_schema.py
-1
View File
@@ -68,7 +68,6 @@ steps:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
@@ -534,6 +534,20 @@ else
echo "--- Single-node job"
echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
ulimit_core_hard=$(ulimit -H -c)
if [[ "$ulimit_core_hard" == "unlimited" ]]; then
# docker run can't pass "unlimited" to --ulimit
ulimit_core_hard="-1"
fi
# Disable core dumps in the ROCm test container unless the ROCm debug agent is enabled
coredump_flags="--ulimit core=0:$ulimit_core_hard"
if [[ "$commands" == *"ROCm debug agent enabled"* ]]; then
# Works around https://github.com/rocm/rocm-systems/issues/6206
coredump_flags='-e HSA_COREDUMP_PATTERN="/tmp/gpucore.%p"'
else
echo "ROCm debug agent not enabled, coredumps are disabled in the test container."
fi
docker run \
--device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
$RDMA_FLAGS \
@@ -541,6 +555,7 @@ else
--shm-size=16gb \
--group-add "$render_gid" \
--rm \
$coredump_flags \
-e HF_TOKEN \
-e "HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT}" \
-e "HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT}" \
@@ -1,20 +0,0 @@
#!/bin/bash
# This script build the CPU docker image and run the offline inference inside the container.
# It serves a sanity check for compilation and basic model usage.
set -euox pipefail
# allow to bind to different cores
CORE_RANGE=${CORE_RANGE:-48-95}
NUMA_NODE=${NUMA_NODE:-1}
IMAGE_NAME="zen-cpu-test-$NUMA_NODE"
TIMEOUT_VAL=$1
TEST_COMMAND=$2
# building the docker image
echo "--- :docker: Building Docker image"
docker build --progress plain --tag "$IMAGE_NAME" --target vllm-zen-test -f docker/Dockerfile.cpu .
# Run the image, setting --shm-size=4g for tensor parallel.
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
+34 -50
View File
@@ -114,26 +114,6 @@ steps:
#---------------------------------------------------------- mi250 · compile ----------------------------------------------------------#
- label: PyTorch Compilation Unit Tests # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
torch_nightly: true
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/compilation/
- vllm/model_executor/layers/
- vllm/v1/worker/
- vllm/v1/attention/
- vllm/v1/cudagraph_dispatcher.py
- vllm/config/compilation.py
- csrc/
- tests/compile
- vllm/platforms/rocm.py
commands:
- "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
- label: PyTorch Fullgraph Smoke Test # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -259,7 +239,6 @@ steps:
- vllm/
- tests/models/multimodal
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"
- pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
@@ -438,7 +417,7 @@ steps:
#----------------------------------------------------- mi300 · basic_correctness -----------------------------------------------------#
- label: Basic Correctness # TBD
timeout_in_minutes: 50
timeout_in_minutes: 95
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
fast_check: true
@@ -456,7 +435,7 @@ steps:
- pytest -v -s basic_correctness/test_cpu_offload.py
- label: Distributed Model Tests (2 GPUs) # TBD
timeout_in_minutes: 65
timeout_in_minutes: 110
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_2
num_gpus: 2
@@ -498,6 +477,26 @@ steps:
#---------------------------------------------------------- mi300 · compile ----------------------------------------------------------#
- label: PyTorch Compilation Unit Tests # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
torch_nightly: true
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/compilation/
- vllm/model_executor/layers/
- vllm/v1/worker/
- vllm/v1/attention/
- vllm/v1/cudagraph_dispatcher.py
- vllm/config/compilation.py
- csrc/
- tests/compile
- vllm/platforms/rocm.py
commands:
- "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
- label: Fusion E2E Config Sweep (H100-MI300) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -678,7 +677,7 @@ steps:
- pytest -v -s distributed/test_eplb_spec_decode.py
- label: Distributed Tests (2xH100-2xMI300) # TBD
timeout_in_minutes: 30
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_2
num_gpus: 2
@@ -1222,7 +1221,7 @@ steps:
#--------------------------------------------------------- mi300 · examples ----------------------------------------------------------#
- label: Examples # TBD
timeout_in_minutes: 45
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
@@ -1258,7 +1257,7 @@ steps:
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
- label: Kernels Attention Test %N # TBD
timeout_in_minutes: 55
timeout_in_minutes: 100
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
@@ -1292,7 +1291,7 @@ steps:
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
- label: Kernels MoE Test %N # TBD
timeout_in_minutes: 50
timeout_in_minutes: 95
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
@@ -1374,8 +1373,6 @@ steps:
- tests/lora
- vllm/platforms/rocm.py
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
- pytest -v -s -x lora/test_chatglm3_tp.py
- pytest -v -s -x lora/test_llama_tp.py
- pytest -v -s -x lora/test_qwen3_with_multi_loras.py
@@ -1439,7 +1436,7 @@ steps:
- pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
- label: Basic Models Tests (Other) # TBD
timeout_in_minutes: 45
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
@@ -1516,7 +1513,6 @@ steps:
- tests/models/multimodal/generation
- tests/models/multimodal/test_mapping.py
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
- pytest -v -s models/multimodal/test_mapping.py
@@ -1530,7 +1526,6 @@ steps:
- vllm/
- tests/models/multimodal/generation
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
@@ -1544,7 +1539,6 @@ steps:
- vllm/
- tests/models/multimodal/generation
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD
@@ -1558,7 +1552,6 @@ steps:
- vllm/
- tests/models/multimodal
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2"
- pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
@@ -1574,7 +1567,6 @@ steps:
- tests/models/multimodal/generation
- tests/models/multimodal/test_mapping.py
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"
- pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
@@ -1589,7 +1581,6 @@ steps:
- tests/models/multimodal/generation
- tests/models/multimodal/test_mapping.py
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing
- pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model
@@ -1605,7 +1596,6 @@ steps:
- tests/models/multimodal
- tests/models/registry.py
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/processing/test_tensor_schema.py
- label: Multi-Modal Processor (CPU) %N # TBD
@@ -1621,7 +1611,6 @@ steps:
- tests/models/multimodal
- tests/models/registry.py
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
#----------------------------------------------------- mi300 · models / quantized -----------------------------------------------------#
@@ -1903,7 +1892,7 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- label: Spec Decode Eagle # TBD
timeout_in_minutes: 45
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
@@ -2119,7 +2108,7 @@ steps:
- DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
- label: Metrics, Tracing (2 GPUs) # TBD
timeout_in_minutes: 20
timeout_in_minutes: 65
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_2
optional: true
@@ -2272,7 +2261,7 @@ steps:
#------------------------------------------------------ mi300 · weight_loading -------------------------------------------------------#
- label: Weight Loading Multiple GPU # TBD
timeout_in_minutes: 30
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_2
num_gpus: 2
@@ -2284,7 +2273,7 @@ steps:
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
- label: Weight Loading Multiple GPU - Large Models # TBD
timeout_in_minutes: 30
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_2
num_gpus: 2
@@ -2469,7 +2458,6 @@ steps:
- vllm/
- tests/models/multimodal
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
@@ -2831,7 +2819,7 @@ steps:
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
- label: Kernels Attention Test %N # TBD
timeout_in_minutes: 60
timeout_in_minutes: 100
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_1
parallelism: 2
@@ -2980,7 +2968,6 @@ steps:
- tests/models/multimodal/generation
- tests/models/multimodal/test_mapping.py
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
- pytest -v -s models/multimodal/test_mapping.py
@@ -2994,7 +2981,6 @@ steps:
- vllm/
- tests/models/multimodal/generation
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
- label: Multi-Modal Models (Extended Pooling) # TBD
@@ -3020,7 +3006,6 @@ steps:
- vllm/
- tests/models/multimodal
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2"
- pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
@@ -3035,7 +3020,6 @@ steps:
- vllm/
- tests/models/multimodal/generation
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing
- pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model
@@ -3178,7 +3162,7 @@ steps:
#------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------#
- label: Weight Loading Multiple GPU # TBD
timeout_in_minutes: 30
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_2
num_gpus: 2
@@ -3190,7 +3174,7 @@ steps:
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
- label: Weight Loading Multiple GPU - Large Models # TBD
timeout_in_minutes: 30
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_2
working_dir: "/vllm-workspace/tests"
+1 -1
View File
@@ -233,7 +233,7 @@ steps:
num_devices: 2
commands:
- pytest -v -s tests/distributed/test_context_parallel.py
- pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
- pytest -v -s tests/distributed/test_nccl_symm_mem.py
- pytest -v -s tests/v1/distributed/test_dbo.py
- pytest -v -s tests/distributed/test_mnnvl_alltoall.py
+6 -5
View File
@@ -54,8 +54,8 @@ steps:
- export VLLM_USE_DEEP_GEMM=0 # We found Triton is faster than DeepGEMM for H100
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
- label: LM Eval Small Models (2xB200)
key: lm-eval-small-models-2xb200
- label: LM Eval Small Models (1xB200)
key: lm-eval-small-models-1xb200
timeout_in_minutes: 120
device: b200-k8s
optional: true
@@ -65,9 +65,10 @@ steps:
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
- label: LM Eval Small Models (2xL4)
key: lm-eval-small-models-tp
timeout_in_minutes: 10
- label: LM Eval Small Models Distributed (2xB200)
key: lm-eval-small-models-distributed-2xb200
timeout_in_minutes: 120
device: b200-k8s
num_devices: 2
optional: true
source_file_dependencies:
@@ -10,7 +10,6 @@ steps:
- vllm/
- tests/models/multimodal
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2"
- pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
mirror:
@@ -27,7 +26,6 @@ steps:
- vllm/
- tests/models/multimodal
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
mirror:
@@ -44,7 +42,6 @@ steps:
- vllm/
- tests/models/multimodal
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"
- pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
mirror:
@@ -61,7 +58,6 @@ steps:
- vllm/
- tests/models/multimodal
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
- pytest models/multimodal/generation/test_memory_leak.py -m core_model
@@ -83,7 +79,6 @@ steps:
- tests/models/registry.py
device: cpu-medium
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
- label: Multi-Modal Processor # 44min
@@ -95,7 +90,6 @@ steps:
- tests/models/multimodal
- tests/models/registry.py
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/processing/test_tensor_schema.py
- label: Multi-Modal Accuracy Eval (Small Models) # 50min
@@ -129,7 +123,6 @@ steps:
- tests/models/multimodal/generation
- tests/models/multimodal/test_mapping.py
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
- pytest -v -s models/multimodal/test_mapping.py
mirror:
@@ -146,7 +139,6 @@ steps:
- vllm/
- tests/models/multimodal/generation
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
- label: Multi-Modal Models (Extended Generation 3)
@@ -157,7 +149,6 @@ steps:
- vllm/
- tests/models/multimodal/generation
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
- label: Multi-Modal Models (Extended Pooling)
+2 -4
View File
@@ -132,10 +132,8 @@ def benchmark_function(
reset_memory_stats()
# Benchmark
start_events = [
torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)
]
end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)]
start_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
end_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
for i in range(benchmark_iters):
logits_copy = logits.clone()
+2 -2
View File
@@ -134,8 +134,8 @@ def benchmark_config(
torch.accelerator.synchronize()
# Benchmark
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start = torch.Event(enable_timing=True)
end = torch.Event(enable_timing=True)
start.record()
for _ in range(num_iters):
with override_config(config):
@@ -170,8 +170,8 @@ def benchmark_config(
graph.replay()
torch.accelerator.synchronize()
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start = torch.Event(enable_timing=True)
end = torch.Event(enable_timing=True)
latencies: list[float] = []
for _ in range(num_iters):
start.record()
+1
View File
@@ -427,6 +427,7 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
"csrc/cpu/activation_lut_bf16.cpp"
"csrc/cpu/cpu_tanhf_neon.hpp"
"csrc/cpu/cpu_fused_moe.cpp"
${VLLM_EXT_SRC})
endif()
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG b3964b1d8b95d8e8447435668ab169a2700bab65
GIT_TAG 2c839c33742309ec41e620bf837495ec9926c56e
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+12
View File
@@ -126,6 +126,18 @@ void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
});
}
void gelu_tanh(torch::Tensor& out, torch::Tensor& input) {
int num_tokens = input.numel() / input.size(-1);
int d = input.size(-1);
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_tanh_impl", [&] {
CPU_KERNEL_GUARD_IN(gelu_tanh_impl)
activation_kernel<scalar_t, gelu_tanh_act, false>(
num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
CPU_KERNEL_GUARD_OUT(gelu_tanh_impl)
});
}
void gelu_new(torch::Tensor& out, torch::Tensor& input) {
int num_tokens = input.numel() / input.size(-1);
int d = input.size(-1);
+128
View File
@@ -0,0 +1,128 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#ifndef CPU_TANHF_NEON_HPP
#define CPU_TANHF_NEON_HPP
#include <cstdint>
#include <arm_neon.h>
namespace vec_op {
namespace {
struct TanhfConstants {
float32x4_t special_bound;
float32x4_t two;
float32x4_t c0;
float32x4_t c2;
int32x4_t exponent_bias;
float c1;
float c3;
float two_over_ln2;
float c4;
float ln2_hi;
float ln2_lo;
};
const TanhfConstants kTanhfConstants = {
// 9.01, above which tanhf rounds to 1 (or -1 for negative).
.special_bound = vdupq_n_f32(0x1.205966p+3f),
.two = vdupq_n_f32(0x1.0p+1f),
.c0 = vdupq_n_f32(0x1.fffffep-2f),
.c2 = vdupq_n_f32(0x1.555736p-5f),
.exponent_bias = vdupq_n_s32(0x3f800000),
.c1 = 0x1.5554aep-3f,
.c3 = 0x1.12287cp-7f,
.two_over_ln2 = 0x1.715476p+1f,
.c4 = 0x1.6b55a2p-10f,
.ln2_hi = 0x1.62e4p-1f,
.ln2_lo = 0x1.7f7d1cp-20f,
};
// Return the ptr but hide it's value from the compiler so accesses
// through it can't be optimised based on contents.
template <typename T>
inline const T* ptr_barrier(const T* ptr) {
const T* opaque_ptr = ptr;
__asm__("" : "+r"(opaque_ptr));
return opaque_ptr;
}
// Check whether any lanes in the mask are set
inline bool any_u32(uint32x4_t x) { return vmaxvq_u32(x) != 0; }
// e^2x - 1 inline helper
inline float32x4_t e2xm1f_inline(float32x4_t x, const TanhfConstants* d) {
float32x2_t ln2 = vld1_f32(&d->ln2_hi);
float32x4_t lane_consts = vld1q_f32(&d->c1);
// Reduce argument: f in [-ln2/2, ln2/2], i is exact.
float32x4_t j = vrndaq_f32(vmulq_laneq_f32(x, lane_consts, 2));
int32x4_t i = vcvtq_s32_f32(j);
float32x4_t f = vaddq_f32(x, x);
f = vfmsq_lane_f32(f, j, ln2, 0);
f = vfmsq_lane_f32(f, j, ln2, 1);
// Approximate expm1(f) with polynomial P, expm1(f) ~= f + f^2 * P(f)
float32x4_t f2 = vmulq_f32(f, f);
float32x4_t f4 = vmulq_f32(f2, f2);
float32x4_t p01 = vfmaq_laneq_f32(d->c0, f, lane_consts, 0);
float32x4_t p23 = vfmaq_laneq_f32(d->c2, f, lane_consts, 1);
float32x4_t poly = vfmaq_f32(p01, f2, p23);
poly = vfmaq_laneq_f32(poly, f4, lane_consts, 3);
poly = vfmaq_f32(f, f2, poly);
// scale = 2^i
int32x4_t u = vaddq_s32(vshlq_n_s32(i, 23), d->exponent_bias);
float32x4_t scale = vreinterpretq_f32_s32(u);
return vfmaq_f32(vsubq_f32(scale, vdupq_n_f32(1.0f)), poly, scale);
}
// Calculate the result tanh(x) = q / (q+2) and set special lanes to ±1
inline float32x4_t special_case(float32x4_t x, float32x4_t q,
uint32x4_t special) {
const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
float32x4_t y = vdivq_f32(q, vaddq_f32(q, d->two));
uint32x4_t ix = vreinterpretq_u32_f32(x);
uint32x4_t one_bits = vreinterpretq_u32_s32(d->exponent_bias);
uint32x4_t sign_mask = vdupq_n_u32(0x80000000u);
uint32x4_t special_bits = vbslq_u32(sign_mask, ix, one_bits);
float32x4_t special_y = vreinterpretq_f32_u32(special_bits);
return vbslq_f32(special, special_y, y);
}
} // namespace
// Implementation of tanhf adapted from Arm Optimized Routines (tanhf
// AdvSIMD)
// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/tanhf.c
//
// Approximation for single-precision vector tanh(x), using a simplified
// version of expm1f. The maximum error is 2.08 + 0.5 ULP:
// _ZGVnN4v_tanhf (0x1.fa5eep-5) got 0x1.f9ba02p-5 want 0x1.f9ba08p-5.
inline float32x4_t fast_tanhf_f32x4(float32x4_t x) {
const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
// tanh(x) = (e^2x - 1) / (e^2x + 1)
// q = e^2x -1
float32x4_t q = e2xm1f_inline(x, d);
// Check for special cases
uint32x4_t special = vcagtq_f32(x, d->special_bound);
// Fall back to vectorised special case for any lanes which would cause
// expm1 to overflow
if (any_u32(special)) {
return special_case(x, q, special);
}
// Complete fast path if no special lanes
// tanh(x) = q / (q+2)
return vdivq_f32(q, vaddq_f32(q, d->two));
}
} // namespace vec_op
#endif // CPU_TANHF_NEON_HPP
+22
View File
@@ -3,6 +3,8 @@
#include <arm_neon.h>
#include "cpu/cpu_tanhf_neon.hpp"
#include <torch/all.h>
#include <ATen/cpu/vec/functional.h>
#include <ATen/cpu/vec/vec.h>
@@ -345,6 +347,10 @@ struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
explicit FP32Vec4(float32x4_t data) : Base(VectorizedT(data)) {};
explicit FP32Vec4(const FP32Vec4& data) : Base(data) {};
FORCE_INLINE FP32Vec4 tanh() const {
return FP32Vec4(fast_tanhf_f32x4(reg.val[0]));
}
};
struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
@@ -391,6 +397,13 @@ struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
reg.val[1] = Vectorized<float>(data.val[1]);
}
FORCE_INLINE FP32Vec8 tanh() const {
FP32Vec8 r(uninit);
r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
return r;
}
FORCE_INLINE float reduce_sum() const noexcept {
float answer = 0;
std::plus<VectorizedT> add;
@@ -497,6 +510,15 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
reg.val[3] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[1])));
};
FORCE_INLINE FP32Vec16 tanh() const {
FP32Vec16 r(uninit);
r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
r.reg.val[2] = Vectorized<float>(fast_tanhf_f32x4(reg.val[2]));
r.reg.val[3] = Vectorized<float>(fast_tanhf_f32x4(reg.val[3]));
return r;
}
static FORCE_INLINE void load_even_odd(const float* ptr, FP32Vec16& even,
FP32Vec16& odd) noexcept {
const float32x4x2_t x01 = vuzpq_f32(vld1q_f32(ptr), vld1q_f32(ptr + 4));
+4
View File
@@ -298,6 +298,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
ops.impl("gelu_tanh_and_mul", torch::kCPU, &gelu_tanh_and_mul);
// GELU tanh implementation.
ops.def("gelu_tanh(Tensor! out, Tensor input) -> ()");
ops.impl("gelu_tanh", torch::kCPU, &gelu_tanh);
// GELU implementation used in GPT-2.
ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
ops.impl("gelu_new", torch::kCPU, &gelu_new);
@@ -804,35 +804,6 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) {
} // namespace tensorrt_llm
} // namespace vllm
torch::stable::Tensor minimax_allreduce_rms(
torch::stable::Tensor const& input,
torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
int64_t const rank, int64_t const nranks, double const eps) {
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
allreduce_params.nranks = static_cast<int>(nranks);
allreduce_params.rank = static_cast<int>(rank);
allreduce_params.dtype = input.scalar_type();
allreduce_params.size_q = static_cast<int>(input.numel());
allreduce_params.hidden_dim = static_cast<int>(input.size(-1));
allreduce_params.stride_q = allreduce_params.hidden_dim;
allreduce_params.workspace =
reinterpret_cast<void**>(workspace.mutable_data_ptr());
allreduce_params.allreduce_in = const_cast<void*>(input.const_data_ptr());
allreduce_params.rms_gamma = const_cast<void*>(norm_weight.const_data_ptr());
allreduce_params.rms_eps = static_cast<float>(eps);
allreduce_params.stream = get_current_cuda_stream(input.get_device_index());
torch::stable::Tensor rms_norm_out = torch::stable::empty_like(input);
allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr();
vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params);
return rms_norm_out;
}
std::tuple<torch::stable::Tensor, torch::stable::Tensor>
minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
torch::stable::Tensor const& norm_weight_q,
-4
View File
@@ -288,10 +288,6 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
int64_t cache_block_size);
#ifndef USE_ROCM
torch::stable::Tensor minimax_allreduce_rms(
torch::stable::Tensor const& input,
torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
int64_t const rank, int64_t const nranks, double const eps);
std::tuple<torch::stable::Tensor, torch::stable::Tensor>
minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
torch::stable::Tensor const& norm_weight_q,
-5
View File
@@ -449,10 +449,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"int cache_block_size) -> ()");
#ifndef USE_ROCM
ops.def(
"minimax_allreduce_rms("
"Tensor input, Tensor norm_weight, Tensor workspace, "
"int rank, int nranks, float eps) -> Tensor");
ops.def(
"minimax_allreduce_rms_qk("
"Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, "
@@ -705,7 +701,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
#ifndef USE_ROCM
ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
#endif
ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
+2
View File
@@ -35,6 +35,8 @@ void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
void gelu_tanh(torch::Tensor& out, torch::Tensor& input);
void gelu_new(torch::Tensor& out, torch::Tensor& input);
void gelu_fast(torch::Tensor& out, torch::Tensor& input);
-27
View File
@@ -318,30 +318,3 @@ RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install "vllm[zen]"
ENTRYPOINT ["vllm", "serve"]
######################### ZEN CPU TEST IMAGE #########################
FROM vllm-openai-zen AS vllm-zen-test
WORKDIR /vllm-workspace
COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -r requirements/test/cpu.txt
ADD ./tests/ ./tests/
ADD ./examples/ ./examples/
ADD ./benchmarks/ ./benchmarks/
ADD ./vllm/collect_env.py .
ADD ./docker/ ./docker/
ADD ./.buildkite/ ./.buildkite/
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -e tests/vllm_test_utils
# enable fast downloads from hf (for testing)
ENV HF_XET_HIGH_PERFORMANCE=1
# increase timeout for hf downloads (for testing)
ENV HF_HUB_DOWNLOAD_TIMEOUT=60
ENTRYPOINT []
+1 -1
View File
@@ -222,7 +222,7 @@ MLA decode backends are selected using the standard
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | | Decoder | 10.x |
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | | Decoder | 10.x |
| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
+3 -8
View File
@@ -405,7 +405,6 @@ th {
| `Glm4MoeLiteForCausalLM` | GLM-4.7-Flash | `zai-org/GLM-4.7-Flash`, etc. | ✅︎ | ✅︎ |
| `GlmMoeDsaForCausalLM` | GLM-5, GLM-5.1, GLM-5.2 | `zai-org/GLM-5`, etc. | ✅︎ | ✅︎ |
| `GPT2LMHeadModel` | GPT-2 | `openai-community/gpt2`, `openai-community/gpt2-xl`, etc. | | ✅︎ |
| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | ✅︎ |
| `GPTJForCausalLM` | GPT-J | `EleutherAI/gpt-j-6b`, `nomic-ai/gpt4all-j`, etc. | | ✅︎ |
| `GPTNeoXForCausalLM` | GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM | `EleutherAI/gpt-neox-20b`, `EleutherAI/pythia-12b`, `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, `databricks/dolly-v2-12b`, `stabilityai/stablelm-tuned-alpha-7b`, etc. | | ✅︎ |
| `GptOssForCausalLM` | GPT-OSS | `openai/gpt-oss-120b`, `openai/gpt-oss-20b` | ✅︎ | ✅︎ |
@@ -477,7 +476,6 @@ th {
| `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
| `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |
| `StableLMEpochForCausalLM` | StableLM Epoch | `stabilityai/stablelm-zephyr-3b`, etc. | | ✅︎ |
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | | ✅︎ |
| `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ |
| `Step3p5ForCausalLM` | Step-3.5-flash | `stepfun-ai/Step-3.5-Flash`, etc. | | ✅︎ |
| `TeleChatForCausalLM` | TeleChat | `chuhac/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
@@ -490,7 +488,9 @@ Some models are supported only via the [Transformers modeling backend](#transfor
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | |
| `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ |
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ |
!!! note
Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096.
@@ -532,7 +532,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| ------------ | ------ | ------ | ----------------- | -------------------- | ------------------------- |
| `AriaForConditionalGeneration` | Aria | T + I<sup>+</sup> | `rhymes-ai/Aria` | | |
| `AudioFlamingo3ForConditionalGeneration` | AudioFlamingo3 | T + A | `nvidia/audio-flamingo-3-hf`, `nvidia/music-flamingo-hf` | ✅︎ | ✅︎ |
| `AyaVisionForConditionalGeneration` | Aya Vision | T + I<sup>+</sup> | `CohereLabs/aya-vision-8b`, `CohereLabs/aya-vision-32b`, etc. | | ✅︎ |
| `BagelForConditionalGeneration` | BAGEL | T + I<sup>+</sup> | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ |
| `BeeForConditionalGeneration` | Bee-8B | T + I<sup>E+</sup> | `Open-Bee/Bee-8B-RL`, `Open-Bee/Bee-8B-SFT` | | ✅︎ |
| `Blip2ForConditionalGeneration` | BLIP-2 | T + I<sup>E</sup> | `Salesforce/blip2-opt-2.7b`, `Salesforce/blip2-opt-6.7b`, etc. | ✅︎ | ✅︎ |
@@ -579,7 +578,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
| `Llama4ForConditionalGeneration` | Llama 4 | T + I<sup>+</sup> | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ |
| `Llama_Nemotron_Nano_VL` | Llama Nemotron Nano VL | T + I<sup>E+</sup> | `nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1` | ✅︎ | ✅︎ |
| `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + I<sup>E+</sup> | `llava-hf/llava-1.5-7b-hf`, `TIGER-Lab/Mantis-8B-siglip-llama3` (see note), `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ |
| `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + I<sup>E+</sup> | `llava-hf/llava-1.5-7b-hf`, `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ |
| `LlavaNextForConditionalGeneration` | LLaVA-NeXT, Granite Vision | T + I<sup>E+</sup> | `llava-hf/llava-v1.6-mistral-7b-hf`, `llava-hf/llava-v1.6-vicuna-7b-hf`, `ibm-granite/granite-vision-3.3-2b`, etc. | | ✅︎ |
| `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |
| `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
@@ -594,7 +593,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`<sup>^</sup>, `allenai/MolmoWeb-8B`<sup>^</sup> | ✅︎ | ✅︎ |
| `MossAudioModel` | MOSS-Audio | T + A<sup>+</sup> | `OpenMOSS-Team/MOSS-Audio-4B-Instruct`, `OpenMOSS-Team/MOSS-Audio-4B-Thinking`, `OpenMOSS-Team/MOSS-Audio-8B-Instruct`, `OpenMOSS-Team/MOSS-Audio-8B-Thinking` | ✅︎ | ✅︎ |
| `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
| `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ |
| `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
| `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + I<sup>E+</sup> | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
| `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ |
@@ -678,9 +676,6 @@ Some models are supported only via the [Transformers modeling backend](#transfor
coordinate decoding and are not exposed by this vLLM implementation.
See [Moondream3 prompt recipes](../features/multimodal_inputs.md#moondream3-prompt-recipes).
!!! note
To use `TIGER-Lab/Mantis-8B-siglip-llama3`, you have to pass `--hf_overrides '{"architectures": ["MantisForConditionalGeneration"]}'` when running vLLM.
!!! note
The official `openbmb/MiniCPM-V-2` doesn't work yet, so we need to use a fork (`HwwwH/MiniCPM-V-2`) for now.
For more details, please see: <https://github.com/vllm-project/vllm/pull/4087#issuecomment-2250397630>
+1 -1
View File
@@ -58,7 +58,7 @@ class Fp8PerTensorOnlineLinearMethod(LinearMethodBase):
### High Level Weight Transfer API
The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Layerwise reloading is controlled by the `WeightTransferUpdateInfo.is_checkpoint_format` flag and is set to `True` by default.
The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Checkpoint-format weight transfer engines (e.g. the NCCL and IPC backends) run layerwise reloading automatically inside their `start_weight_update`/`finish_weight_update` lifecycle.
### Mid Level `reload_weights` API
+3 -2
View File
@@ -17,6 +17,7 @@ The weight transfer system follows a **four-phase protocol** with a pluggable ba
| ------- | --------- | -------- |
| [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference |
| [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) |
## Configuration
@@ -41,7 +42,7 @@ vllm serve my-model \
--weight-transfer-config '{"backend": "nccl"}'
```
The `backend` field accepts `"nccl"` (default) or `"ipc"`.
The `backend` field accepts `"nccl"` (default), `"ipc"`, or `"sparse_nccl"`.
## API Endpoints
@@ -69,7 +70,7 @@ Both backends provide static methods that the trainer calls to send weights. The
EngineClass.trainer_init(init_info)
# 2. Start weight update on inference side
llm.start_weight_update(is_checkpoint_format=True)
llm.start_weight_update()
# 3. Send weights to inference workers
EngineClass.trainer_send_weights(
+28 -15
View File
@@ -11,15 +11,23 @@ The `WeightTransferEngine` is a generic abstract class parameterized by two data
### Abstract Methods
Subclasses must implement these four methods:
Subclasses must implement these methods:
| Method | Side | Description |
| ------ | ---- | ----------- |
| `init_transfer_engine(init_info)` | Inference | Initialize the communication channel on each inference worker |
| `receive_weights(update_info, load_weights)` | Inference | Receive weights and call `load_weights` incrementally |
| `start_weight_update()` | Inference | Prepare for an update (e.g. begin layerwise reload); no-op for in-place engines |
| `finish_weight_update()` | Inference | Finalize the update (e.g. finalize layerwise reload); no-op for in-place engines |
| `receive_weights(update_info)` | Inference | Receive weights and load them into `self.model` |
| `shutdown()` | Inference | Clean up resources |
| `trainer_send_weights(iterator, trainer_args)` | Trainer | Static method to send weights from the trainer process |
The base class provides two methods:
1. `__init__` : Engines receive `config` (`WeightTransferConfig`), `vllm_config` (`VllmConfig`), `device` (`torch.device`) and `model` (`nn.Module`)
2. `update_weights(update_info_dict)`: Thin wrapper for `receive_weights`: parses
the dict into user-specified data type, calls `receive_weights`, and synchronizes the device. Subclasses implement `receive_weights`.
### Request Classes
The API-level request classes provide backend-agnostic serialization using plain dictionaries. The engine's `parse_init_info` and `parse_update_info` methods convert these dictionaries into typed dataclasses.
@@ -81,7 +89,7 @@ class MyUpdateInfo(WeightTransferUpdateInfo):
### 2. Implement the Engine
```python
from collections.abc import Callable, Iterator
from collections.abc import Iterator
from typing import Any
import torch
@@ -93,18 +101,25 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
# Set up connection to trainer using init_info.endpoint, etc.
...
def receive_weights(
self,
update_info: MyUpdateInfo,
load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
) -> None:
# Receive each weight and call load_weights incrementally
def start_weight_update(self) -> None:
# Checkpoint-format engines: run initialize_layerwise_reload(self.model).
# In-place engines: no-op
...
def finish_weight_update(self) -> None:
# Checkpoint-format engines: run finalize_layerwise_reload(...).
# In-place engines: no-op
...
def receive_weights(self, update_info: MyUpdateInfo) -> None:
weights = []
for name, dtype_name, shape in zip(
update_info.names, update_info.dtype_names, update_info.shapes
):
dtype = getattr(torch, dtype_name)
weight = self._fetch_weight(name, shape, dtype)
load_weights([(name, weight)])
weights.append((name, weight))
self.model.load_weights(weights)
def shutdown(self) -> None:
# Clean up resources
@@ -121,9 +136,6 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
...
```
!!! important
The `load_weights` callable passed to `receive_weights` should be called **incrementally** (one or a few weights at a time) rather than accumulating all weights first. This avoids GPU out-of-memory errors with large models.
### 3. Register with the Factory
```python
@@ -147,7 +159,7 @@ Once registered, users can select your backend via `WeightTransferConfig(backend
## WeightTransferEngineFactory
The factory uses a registry pattern with lazy loading. Built-in engines (`nccl` and `ipc`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed.
The factory uses a registry pattern with lazy loading. Built-in engines (`nccl`, `ipc`, and `sparse_nccl`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed.
```python
from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
@@ -155,7 +167,8 @@ from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
# Create an engine from config
engine = WeightTransferEngineFactory.create_engine(
config=weight_transfer_config,
parallel_config=parallel_config,
vllm_config=vllm_config,
device=device,
model=model,
)
```
+2 -2
View File
@@ -55,7 +55,7 @@ trainer_args = IPCTrainerSendWeightsArgs(
llm_handle=llm_actor_handle,
)
# start
ray.get(llm_actor_handle.start_weight_update.remote(is_checkpoint_format=True))
ray.get(llm_actor_handle.start_weight_update.remote())
# send weights
IPCWeightTransferEngine.trainer_send_weights(
iterator=model.named_parameters(),
@@ -80,7 +80,7 @@ trainer_args = IPCTrainerSendWeightsArgs(
# start
base_url = "http://localhost:8000"
url = f"{base_url}/start_weight_update"
response = requests.post(url, json={"is_checkpoint_format": True}, timeout=60)
response = requests.post(url, json={}, timeout=60)
response.raise_for_status()
# send weights
IPCWeightTransferEngine.trainer_send_weights(
+14 -11
View File
@@ -11,7 +11,7 @@ The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) b
## How It Works
1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction).
2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads weights incrementally.
2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads the weights.
3. Optionally, **packed tensor broadcasting** batches multiple small tensors into larger buffers with double/triple buffering and CUDA stream overlap for higher throughput. This implementation is based on [NeMo-RL's packed tensor](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/utils/packed_tensor.py).
## Initialization
@@ -93,7 +93,7 @@ remaining three steps are:
from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
# 1. Start the weight update
llm.start_weight_update(is_checkpoint_format=True)
llm.start_weight_update()
# 2. Receive weights (can be called multiple times for chunked transfers)
llm.update_weights(
@@ -116,19 +116,22 @@ must match the order in which the trainer iterates over its parameters.
`start_weight_update` must be called before `update_weights`, and
`finish_weight_update` must be called after all weight chunks have been
transferred. The `is_checkpoint_format` flag controls whether layerwise reload
processing is applied (`True` for checkpoint-format weights, `False` for
pre-processed kernel-format weights).
transferred. The NCCL engine receives checkpoint-format weights and applies
layerwise reload processing automatically inside `start_weight_update` /
`finish_weight_update`.
Sparse NCCL patches still use `update_kind="sparse_flat"` inside
`update_info`, but they should be wrapped in
`start_weight_update(is_checkpoint_format=False)` because sparse patches apply
directly to runtime/kernel-format parameters. The current sparse MVP requires
`TP=1` and `PP=1`.
## Sparse NCCL
Sparse, flat-index weight patches use a separate backend,
`WeightTransferConfig(backend="sparse_nccl")`, implemented by
`SparseNCCLWeightTransferEngine`. It shares only NCCL process-group
initialization with the dense engine; patches are applied directly in place to
existing parameters (no layerwise reload). The current sparse MVP requires
`TP=1` and `PP=1`. See the example below.
## Examples
- [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1`
- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `backend="sparse_nccl"` and currently require `TP=1` and `PP=1`
- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
- [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
+21
View File
@@ -326,6 +326,27 @@ vLLM supports dynamically loading and unloading LoRA adapters at runtime via the
**Warning:** Dynamic LoRA loading is not a secure operation and should not be enabled in deployments exposed to untrusted clients. If you must enable dynamic LoRA loading, restrict access to the `/v1/load_lora_adapter` and `/v1/unload_lora_adapter` endpoints to trusted administrators only, using a reverse proxy or network-level access controls. Do not expose these endpoints to end users. For details on configuring LoRA adapters, see the [LoRA Adapters documentation](../features/lora.md).
## gRPC Interface
vLLM provides an optional gRPC Generate service on a separate TCP port, enabled via the `--grpc-port` flag. When not specified, no gRPC server is started. The gRPC listener binds to the same host address as the HTTP server.
**Warning:** The gRPC interface is **insecure by default** — it does not implement authentication, authorization, or encryption. It should be considered a private, internal interface intended for use only between co-located services within a trusted network. Do not expose the gRPC port to the public internet or untrusted clients. If you enable the gRPC interface, protect it via network-level access controls such as firewall rules, network segmentation, or deployment on an isolated private network.
### Security Implications
An attacker who can reach the gRPC port can:
1. **Run arbitrary inference** via the `Generate` and `GenerateStream` RPCs without any credentials
2. **Consume GPU and compute resources** by submitting unbounded generation requests
3. **Cause Denial of Service** by exploiting bugs in the gRPC interface that can crash vLLM.
### Recommendations
- Only enable `--grpc-port` when you have a specific need for gRPC-based inference
- Ensure the gRPC port is only accessible from trusted hosts or services
- Use firewall rules to block external access to the gRPC port
- Consider deploying the gRPC interface on a dedicated internal network interface
## Cache Directory Security
vLLM assumes that its cache directories are **private and trusted**. Cache contents are loaded without cryptographic integrity verification, including formats that support arbitrary code execution. If an untrusted user or process can write to vLLM's cache directories, they may be able to crash vLLM or cause it to execute arbitrary code.
@@ -91,44 +91,6 @@ def run_cohere_asr(question: str, audio_count: int) -> ModelRequestData:
)
# MusicFlamingo
def run_musicflamingo(question: str, audio_count: int) -> ModelRequestData:
model_name = "nvidia/music-flamingo-2601-hf"
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
max_num_seqs=2,
limit_mm_per_prompt={"audio": audio_count},
enforce_eager=True,
)
# MusicFlamingo prompt placeholders use <sound>; vLLM's MusicFlamingo
# multimodal processor expands each one into <|sound_bos|> + audio tokens +
# <|sound_eos|> based on extracted audio feature lengths.
audio_placeholder = "<sound>" * audio_count
system_prompt = (
"You are Music Flamingo, a multimodal assistant for language and music. "
"On each turn you receive an audio clip which contains music and optional "
"text, you will receive at least one or both; use your world knowledge and "
"reasoning to help the user with any task. Interpret the entirety of the "
"content any input music--regardlenss of whether the user calls it audio, "
"music, or sound."
)
prompt = (
"<|im_start|>system\n"
f"{system_prompt}<|im_end|>\n"
"<|im_start|>user\n"
f"{audio_placeholder}{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
return ModelRequestData(
engine_args=engine_args,
prompt=prompt,
)
# Gemma3N
def run_gemma3n(question: str, audio_count: int) -> ModelRequestData:
model_name = "google/gemma-3n-E2B-it"
@@ -565,7 +527,6 @@ model_example_map = {
"kimi_audio": run_kimi_audio,
"midashenglm": run_midashenglm,
"minicpmo": run_minicpmo,
"musicflamingo": run_musicflamingo,
"phi4_mm": run_phi4mm,
"qwen2_audio": run_qwen2_audio,
"qwen2_5_omni": run_qwen2_5_omni,
@@ -74,39 +74,6 @@ def load_aria(question: str, image_urls: list[str]) -> ModelRequestData:
)
def load_aya_vision(question: str, image_urls: list[str]) -> ModelRequestData:
model_name = "CohereLabs/aya-vision-8b"
engine_args = EngineArgs(
model=model_name,
max_num_seqs=2,
limit_mm_per_prompt={"image": len(image_urls)},
)
placeholders = [{"type": "image", "image": url} for url in image_urls]
messages = [
{
"role": "user",
"content": [
*placeholders,
{"type": "text", "text": question},
],
}
]
processor = AutoProcessor.from_pretrained(model_name)
prompt = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
return ModelRequestData(
engine_args=engine_args,
prompt=prompt,
image_data=[fetch_image(url) for url in image_urls],
)
def load_bee(question: str, image_urls: list[str]) -> ModelRequestData:
model_name = "Open-Bee/Bee-8B-RL"
@@ -1420,7 +1387,6 @@ def load_molmo2(question: str, image_urls: list[str]) -> ModelRequestData:
model_example_map = {
"aria": load_aria,
"aya_vision": load_aya_vision,
"bee": load_bee,
"command_a_vision": load_command_a_vision,
"deepseek_vl_v2": load_deepseek_vl2,
@@ -68,28 +68,6 @@ def run_aria(questions: list[str], modality: str) -> ModelRequestData:
)
# Aya Vision
def run_aya_vision(questions: list[str], modality: str) -> ModelRequestData:
assert modality == "image"
model_name = "CohereLabs/aya-vision-8b"
engine_args = EngineArgs(
model=model_name,
max_model_len=2048,
max_num_seqs=2,
mm_processor_kwargs={"crop_to_patches": True},
limit_mm_per_prompt={modality: 1},
)
prompts = [
f"<|START_OF_TURN_TOKEN|><|USER_TOKEN|><image>{question}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>"
for question in questions
]
return ModelRequestData(
engine_args=engine_args,
prompts=prompts,
)
# Bee-8B
def run_bee(questions: list[str], modality: str) -> ModelRequestData:
assert modality == "image"
@@ -1377,28 +1355,6 @@ def run_llava_onevision(questions: list[str], modality: str) -> ModelRequestData
)
# Mantis
def run_mantis(questions: list[str], modality: str) -> ModelRequestData:
assert modality == "image"
llama3_template = "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" # noqa: E501
prompts = [llama3_template.format(f"{question}\n<image>") for question in questions]
engine_args = EngineArgs(
model="TIGER-Lab/Mantis-8B-siglip-llama3",
max_model_len=4096,
hf_overrides={"architectures": ["MantisForConditionalGeneration"]},
limit_mm_per_prompt={modality: 1},
)
stop_token_ids = [128009]
return ModelRequestData(
engine_args=engine_args,
prompts=prompts,
stop_token_ids=stop_token_ids,
)
# MiniCPM-V
def run_minicpmv_base(questions: list[str], modality: str, model_name):
assert modality in ["image", "video", "image+video"]
@@ -2349,7 +2305,6 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData:
model_example_map = {
"aria": run_aria,
"aya_vision": run_aya_vision,
"bagel": run_bagel,
"cheers": run_cheers,
"bee": run_bee,
@@ -2390,7 +2345,6 @@ model_example_map = {
"llava-next": run_llava_next,
"llava-next-video": run_llava_next_video,
"llava-onevision": run_llava_onevision,
"mantis": run_mantis,
"minicpmo": run_minicpmo,
"minicpmv": run_minicpmv,
"mistral3": run_mistral3,
+1 -1
View File
@@ -306,7 +306,7 @@ gen_futures = [
ray.get(llm.pause_after_n_tokens.remote())
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
ray.get(llm.start_weight_update.remote())
inference_handle = llm.update_weights.remote(
WeightTransferUpdateRequest(
+3 -7
View File
@@ -80,14 +80,10 @@ def init_weight_transfer_engine(base_url: str) -> None:
response.raise_for_status()
def start_weight_update(
base_url: str,
is_checkpoint_format: bool = True,
) -> None:
def start_weight_update(base_url: str) -> None:
"""Start a weight update via HTTP endpoint."""
url = f"{base_url}/start_weight_update"
payload = {"is_checkpoint_format": is_checkpoint_format}
response = requests.post(url, json=payload, timeout=60)
response = requests.post(url, json={}, timeout=60)
response.raise_for_status()
@@ -170,7 +166,7 @@ def main():
pause_generation(BASE_URL)
# Start weight update, broadcast via IPC, then finish
start_weight_update(BASE_URL, is_checkpoint_format=False)
start_weight_update(BASE_URL)
print("Broadcasting weights via CUDA IPC (HTTP)...")
trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL)
+3 -7
View File
@@ -83,14 +83,10 @@ def init_weight_transfer_engine(
response.raise_for_status()
def start_weight_update(
base_url: str,
is_checkpoint_format: bool = True,
) -> None:
def start_weight_update(base_url: str) -> None:
"""Start a weight update via HTTP endpoint."""
url = f"{base_url}/start_weight_update"
payload = {"is_checkpoint_format": is_checkpoint_format}
response = requests.post(url, json=payload, timeout=60)
response = requests.post(url, json={}, timeout=60)
response.raise_for_status()
@@ -223,7 +219,7 @@ def main():
shapes.append(list(p.shape))
# Start weight update
start_weight_update(BASE_URL, is_checkpoint_format=True)
start_weight_update(BASE_URL)
# Start the update_weights call in a separate thread since it will block
# waiting for NCCL broadcasts
+1 -1
View File
@@ -139,7 +139,7 @@ ray.get(llm.sleep.remote(level=0))
ray.get(train_model.init_weight_transfer.remote())
# Start weight update, sync weights, then finish
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
ray.get(llm.start_weight_update.remote())
ray.get(train_model.broadcast_weights.remote(llm))
ray.get(llm.finish_weight_update.remote())
+3 -10
View File
@@ -277,15 +277,8 @@ class DataParallelInferenceEngine:
]
)
def start_weight_update(self, is_checkpoint_format: bool = True):
ray.get(
[
actor.start_weight_update.remote(
is_checkpoint_format=is_checkpoint_format
)
for actor in self.llm_actors
]
)
def start_weight_update(self):
ray.get([actor.start_weight_update.remote() for actor in self.llm_actors])
def finish_weight_update(self):
ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors])
@@ -392,7 +385,7 @@ def main():
ray.get(inference_engine.wake_up.remote(tags=["weights"]))
print("[sync] Starting weight update...")
ray.get(inference_engine.start_weight_update.remote(is_checkpoint_format=True))
ray.get(inference_engine.start_weight_update.remote())
print("[sync] Packed IPC transfer FSDP → vLLM...")
ray.get(
+1 -1
View File
@@ -202,7 +202,7 @@ ray.get([train_handle, inference_handle])
names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())
# Start weight update
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
ray.get(llm.start_weight_update.remote())
# Issue update_weights call with NCCL-specific update info
# packed=True enables efficient batched tensor broadcasting
+1 -1
View File
@@ -299,7 +299,7 @@ async def main():
print(f"[sync] Got metadata for {len(names)} parameters.")
print("[sync] Starting weight update...")
await engine.start_weight_update(is_checkpoint_format=True)
await engine.start_weight_update()
print("[sync] Broadcasting weights from FSDP → vLLM...")
broadcast_handles = [
+11 -8
View File
@@ -44,11 +44,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
from vllm import LLM, SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.weight_transfer.base import SparseWeightPatch
from vllm.distributed.weight_transfer.nccl_engine import (
NCCLTrainerSendWeightsArgs,
NCCLWeightTransferEngine,
)
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
SparseNCCLWeightTransferEngine,
SparseWeightPatch,
)
from vllm.utils.network_utils import get_ip, get_open_port
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
@@ -244,7 +247,6 @@ class TrainModel:
dtype_names=[str(self.patched_param.dtype).split(".")[-1]],
shapes=[list(self.patched_param.shape)],
num_updates_list=[flat_indices.numel()],
update_kind="sparse_flat",
)
return update_info, selected_token_ids, patch_digest, sparse_payload_bytes
@@ -271,7 +273,7 @@ class TrainModel:
raise RuntimeError("Sparse patch has not been prepared")
start = time.perf_counter()
NCCLWeightTransferEngine.trainer_send_sparse_weights(
SparseNCCLWeightTransferEngine.trainer_send_weights(
iter(self.pending_sparse_patches),
NCCLTrainerSendWeightsArgs(group=self.model_update_group),
)
@@ -282,6 +284,7 @@ class TrainModel:
def launch_llm(
scheduling_inference: PlacementGroupSchedulingStrategy,
backend: str = "nccl",
):
return ray.remote(
num_cpus=0,
@@ -293,7 +296,7 @@ def launch_llm(
tensor_parallel_size=1,
distributed_executor_backend="ray",
gpu_memory_utilization=0.7,
weight_transfer_config=WeightTransferConfig(backend="nccl"),
weight_transfer_config=WeightTransferConfig(backend=backend),
)
@@ -332,7 +335,7 @@ def run_dense_phase(
scheduling_inference: PlacementGroupSchedulingStrategy,
) -> dict[str, object]:
ray.get(train_model.reset_model.remote())
llm = launch_llm(scheduling_inference)
llm = launch_llm(scheduling_inference, backend="nccl")
try:
dense_before = collect_vllm_generations(llm)
@@ -351,7 +354,7 @@ def run_dense_phase(
)
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
ray.get([trainer_init, inference_init])
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
ray.get(llm.start_weight_update.remote())
dense_update_info, dense_payload_bytes = ray.get(
train_model.get_dense_update_info.remote()
@@ -391,7 +394,7 @@ def run_sparse_phase(
scheduling_inference: PlacementGroupSchedulingStrategy,
) -> dict[str, object]:
ray.get(train_model.reset_model.remote())
llm = launch_llm(scheduling_inference)
llm = launch_llm(scheduling_inference, backend="sparse_nccl")
try:
sparse_before = collect_vllm_generations(llm)
@@ -410,7 +413,7 @@ def run_sparse_phase(
)
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
ray.get([trainer_init, inference_init])
ray.get(llm.start_weight_update.remote(is_checkpoint_format=False))
ray.get(llm.start_weight_update.remote())
sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = (
ray.get(train_model.prepare_sparse_patch.remote(PROMPTS))
+16 -1
View File
@@ -5099,6 +5099,7 @@ dependencies = [
"uuid",
"vllm-engine-core-client",
"vllm-llm",
"vllm-model-files",
"vllm-parser",
"vllm-text",
"vllm-tokenizer",
@@ -5236,6 +5237,20 @@ dependencies = [
"zeromq",
]
[[package]]
name = "vllm-model-files"
version = "0.1.0"
dependencies = [
"anyhow",
"hf-hub",
"serde",
"serde_json",
"tempfile",
"thiserror 2.0.18",
"thiserror-ext",
"tokio",
]
[[package]]
name = "vllm-parser"
version = "0.1.0"
@@ -5324,7 +5339,6 @@ dependencies = [
"enum-as-inner",
"expect-test",
"futures",
"hf-hub",
"itertools 0.14.0",
"reqwest",
"serde",
@@ -5339,6 +5353,7 @@ dependencies = [
"trait-set",
"vllm-engine-core-client",
"vllm-llm",
"vllm-model-files",
"vllm-tokenizer",
]
+2
View File
@@ -6,6 +6,7 @@ members = [
"src/llm",
"src/managed-engine",
"src/metrics",
"src/model-files",
"src/mock-engine",
"src/parser",
"src/parser/python",
@@ -126,6 +127,7 @@ vllm-chat = { path = "src/chat" }
vllm-engine-core-client = { path = "src/engine-core-client" }
vllm-llm = { path = "src/llm" }
vllm-managed-engine = { path = "src/managed-engine" }
vllm-model-files = { path = "src/model-files" }
vllm-metrics = { path = "src/metrics" }
vllm-parser = { path = "src/parser" }
vllm-server = { path = "src/server" }
+1
View File
@@ -32,6 +32,7 @@ trait-set.workspace = true
uuid.workspace = true
vllm-engine-core-client.workspace = true
vllm-llm.workspace = true
vllm-model-files.workspace = true
vllm-parser.workspace = true
vllm-text.workspace = true
vllm-tokenizer.workspace = true
+3
View File
@@ -1,5 +1,6 @@
use thiserror::Error;
use thiserror_ext::Macro;
use vllm_model_files::Error as ModelFilesError;
type BoxedError = Box<dyn std::error::Error + Send + Sync>;
@@ -69,6 +70,8 @@ pub enum Error {
#[error(transparent)]
Text(#[from] vllm_text::Error),
#[error(transparent)]
ModelFiles(#[from] ModelFilesError),
#[error(transparent)]
Tokenizer(#[from] vllm_tokenizer::TokenizerError),
}
+6 -6
View File
@@ -6,7 +6,7 @@ use thiserror_ext::AsReport as _;
use tracing::{info, trace, warn};
use vllm_text::Prompt;
use vllm_text::backend::hf::{
HfSpecialTokens, HfTokenizerConfig, ResolvedModelFiles, load_tokenizer_config,
ResolvedModelFiles, SpecialTokens, TokenizerConfig, load_tokenizer_config,
};
use self::format::{
@@ -42,7 +42,7 @@ pub struct HfChatRenderer {
default_template: Option<CompiledChatTemplate>,
default_template_kwargs: HashMap<String, JsonValue>,
content_format: ContentFormatOption,
special_tokens: Option<HfSpecialTokens>,
special_tokens: Option<SpecialTokens>,
multimodal: Option<MultimodalRenderInfo>,
}
@@ -67,7 +67,7 @@ impl HfChatRenderer {
})
}
pub fn with_special_tokens(mut self, special_tokens: Option<HfSpecialTokens>) -> Self {
pub fn with_special_tokens(mut self, special_tokens: Option<SpecialTokens>) -> Self {
self.special_tokens = special_tokens;
self
}
@@ -83,7 +83,7 @@ impl HfChatRenderer {
options: LoadModelBackendsOptions,
multimodal: Option<MultimodalRenderInfo>,
) -> Result<Self> {
let HfTokenizerConfig {
let TokenizerConfig {
special_tokens,
chat_template,
..
@@ -451,7 +451,7 @@ mod tests {
use expect_test::expect;
use serde_json::Value;
use vllm_text::Prompt;
use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken};
use vllm_text::backend::hf::{NamedSpecialToken, SpecialTokens};
use super::{ChatTemplateContentFormatOption, HfChatRenderer, MultimodalRenderInfo};
use crate::request::{
@@ -675,7 +675,7 @@ mod tests {
#[test]
fn chat_template_injects_special_tokens_into_context() {
let request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
let special_tokens = HfSpecialTokens {
let special_tokens = SpecialTokens {
bos_token: Some(NamedSpecialToken::Text("<bos>".to_string())),
..Default::default()
};
+6 -6
View File
@@ -2,7 +2,7 @@
//!
//! This module is inlined from SMG's tokenizer crate with local adaptations:
//! - thinking-related detection/state is removed
//! - special tokens are wired to `vllm_text::backends::hf::HfSpecialTokens`
//! - special tokens are wired to `vllm_text::backends::hf::SpecialTokens`
use std::collections::HashMap;
use std::fs;
@@ -11,7 +11,7 @@ use std::path::Path;
use minijinja::Environment;
use serde::{Deserialize, Serialize};
use serde_json::{self};
use vllm_text::backend::hf::HfSpecialTokens;
use vllm_text::backend::hf::SpecialTokens;
use super::error::TemplateError;
use super::format::{
@@ -46,7 +46,7 @@ pub(super) struct TemplateContext<'a> {
pub(super) tools: Option<&'a [TemplateTool]>,
pub(super) documents: Option<&'a [serde_json::Value]>,
#[serde(flatten)]
pub(super) special_tokens: Option<&'a HfSpecialTokens>,
pub(super) special_tokens: Option<&'a SpecialTokens>,
#[serde(flatten)]
pub(super) template_kwargs: Option<&'a HashMap<String, serde_json::Value>>,
}
@@ -133,7 +133,7 @@ mod tests {
use std::fs;
use tempfile::TempDir;
use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken};
use vllm_text::backend::hf::{NamedSpecialToken, SpecialTokens};
use super::*;
@@ -170,7 +170,7 @@ mod tests {
CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto)
.unwrap();
let special_tokens = HfSpecialTokens {
let special_tokens = SpecialTokens {
bos_token: Some(NamedSpecialToken::Text("<s>".to_string())),
eos_token: Some(NamedSpecialToken::Text("</s>".to_string())),
..Default::default()
@@ -205,7 +205,7 @@ mod tests {
CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto)
.unwrap();
let special_tokens = HfSpecialTokens {
let special_tokens = SpecialTokens {
bos_token: Some(NamedSpecialToken::Text("<s>".to_string())),
eos_token: None,
..Default::default()
+20
View File
@@ -0,0 +1,20 @@
[package]
name = "vllm-model-files"
version.workspace = true
edition.workspace = true
license.workspace = true
[dependencies]
anyhow.workspace = true
hf-hub.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
thiserror-ext.workspace = true
[dev-dependencies]
tempfile.workspace = true
tokio.workspace = true
[lints]
workspace = true
+21
View File
@@ -0,0 +1,21 @@
use std::path::Path;
use serde::Deserialize;
use crate::error::Result;
use crate::json::read_json_file;
/// Minimal subset of `tokenizer_config.json` needed by tokenizer selection.
#[derive(Debug, Default, Deserialize)]
#[serde(default)]
pub(crate) struct TokenizerConfig {
/// The `tokenizer_class` field from HuggingFace tokenizer configs. Some
/// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value
/// containing "Tiktoken" which can be used as a hint for backend
/// selection.
pub tokenizer_class: Option<String>,
}
pub(crate) fn load_tokenizer_config(path: Option<&Path>) -> Result<TokenizerConfig> {
read_json_file(path)
}
+15
View File
@@ -0,0 +1,15 @@
use thiserror::Error as ThisError;
/// Error returned while resolving or reading model files.
#[derive(Debug, ThisError)]
#[error("model file error: {0}")]
pub struct Error(String);
impl Error {
pub(crate) fn new(message: impl Into<String>) -> Self {
Self(message.into())
}
}
/// Result type used by model-file discovery helpers.
pub type Result<T> = std::result::Result<T, Error>;
+31
View File
@@ -0,0 +1,31 @@
use std::fs;
use std::path::Path;
use serde::Deserialize;
use thiserror_ext::AsReport as _;
use crate::error::{Error, Result};
/// Read an optional JSON file into `T`, returning `T::default()` when absent.
pub fn read_json_file<T>(path: Option<&Path>) -> Result<T>
where
T: for<'de> Deserialize<'de> + Default,
{
let Some(path) = path else {
return Ok(T::default());
};
let content = fs::read_to_string(path).map_err(|error| {
Error::new(format!(
"failed to read {}: {}",
path.display(),
error.as_report()
))
})?;
serde_json::from_str(&content).map_err(|error| {
Error::new(format!(
"failed to parse {}: {}",
path.display(),
error.as_report()
))
})
}
+10
View File
@@ -0,0 +1,10 @@
//! Hugging Face model file discovery shared by Rust frontend crates.
mod config;
mod error;
mod json;
mod model_files;
pub use error::{Error, Result};
pub use json::read_json_file;
pub use model_files::{ResolvedModelFiles, TokenizerSource};
@@ -4,7 +4,7 @@ use hf_hub::Cache;
use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo};
use thiserror_ext::AsReport as _;
use super::config::{HfTokenizerConfig, load_tokenizer_config};
use crate::config::{TokenizerConfig, load_tokenizer_config};
use crate::error::{Error, Result};
const HF_TOKEN_ENV: &str = "HF_TOKEN";
@@ -26,6 +26,28 @@ pub enum TokenizerSource {
}
impl TokenizerSource {
/// Select a tokenizer source from a tokenizer file path.
pub fn from_path(path: impl Into<PathBuf>) -> Result<Self> {
let path = path.into();
let file_name = path.file_name().and_then(|name| name.to_str()).ok_or_else(|| {
Error::new(format!(
"tokenizer path has no file name: {}",
path.display()
))
})?;
match file_name {
"tekken.json" => Ok(Self::Tekken(path)),
"tokenizer.json" => Ok(Self::HuggingFace(path)),
_ if is_tiktoken_file(&path) => Ok(Self::Tiktoken(path)),
_ => Err(Error::new(format!(
"unsupported tokenizer file '{}'",
path.display()
))),
}
}
/// Return the local filesystem path for this tokenizer source.
pub fn path(&self) -> &Path {
match self {
Self::HuggingFace(path) | Self::Tiktoken(path) | Self::Tekken(path) => path,
@@ -38,10 +60,15 @@ impl TokenizerSource {
pub struct ResolvedModelFiles {
/// The selected tokenizer source for this model.
pub tokenizer: TokenizerSource,
/// Path to `tokenizer_config.json` when present.
pub tokenizer_config_path: Option<PathBuf>,
/// Path to `generation_config.json` when present.
pub generation_config_path: Option<PathBuf>,
/// Path to `preprocessor_config.json` when present.
pub preprocessor_config_path: Option<PathBuf>,
/// Path to a discovered chat template file when present.
pub chat_template_path: Option<PathBuf>,
/// Path to `config.json` when present.
pub config_path: Option<PathBuf>,
}
@@ -76,10 +103,10 @@ fn resolve_local_model_files(model_dir: &Path) -> Result<ResolvedModelFiles> {
}
async fn resolve_remote_model_files(model_id: &str) -> Result<ResolvedModelFiles> {
let api = build_api().map_err(|error| Error::Tokenizer(error.to_report_string()))?;
let api = build_api().map_err(|error| Error::new(format!("{}", error.as_report())))?;
let repo = api.model(model_id.to_string());
let info = repo.info().await.map_err(|error| {
Error::Tokenizer(format!(
Error::new(format!(
"failed to fetch model '{model_id}': {}",
error.as_report()
))
@@ -138,9 +165,10 @@ fn resolve_cached_model_files(model_id: &str) -> Result<Option<ResolvedModelFile
None => return Ok(None),
};
let model_dir = tokenizer.path().parent().ok_or_else(|| {
Error::Tokenizer("resolved tokenizer file has no parent directory".to_string())
})?;
let model_dir = tokenizer
.path()
.parent()
.ok_or_else(|| Error::new("resolved tokenizer file has no parent directory"))?;
let generation_config_path = cache_repo.get("generation_config.json");
let preprocessor_config_path = cache_repo.get("preprocessor_config.json");
let chat_template_path = discover_chat_template_in_dir(model_dir);
@@ -162,113 +190,88 @@ async fn resolve_remote_tokenizer_source(
siblings: &std::collections::BTreeSet<&str>,
tokenizer_class: Option<&str>,
) -> Result<TokenizerSource> {
if let Some(tekken_path) = download_if_present(repo, model_id, siblings, "tekken.json").await? {
return Ok(TokenizerSource::Tekken(tekken_path));
}
let tokenizer_path = if siblings.contains("tokenizer.json") {
let tokenizer_path = if siblings.contains("tekken.json") {
download_known_file(repo, model_id, "tekken.json").await?
} else if siblings.contains("tokenizer.json") {
download_known_file(repo, model_id, "tokenizer.json").await?
} else if let Some(tiktoken_name) = find_tiktoken_sibling(siblings) {
download_known_file(repo, model_id, tiktoken_name).await?
} else {
return Err(Error::Tokenizer(format!(
return Err(Error::new(format!(
"model '{model_id}' does not expose a supported tokenizer file \
(tokenizer.json, tiktoken.model, or *.tiktoken) on Hugging Face"
)));
};
Ok(resolve_tokenizer_source(
tokenizer_path,
tokenizer_class,
None,
))
resolve_tokenizer_source(tokenizer_path, tokenizer_class)
}
fn resolve_cached_tokenizer_source(
cache_repo: &hf_hub::CacheRepo,
tokenizer_config: &HfTokenizerConfig,
tokenizer_config: &TokenizerConfig,
) -> Result<Option<TokenizerSource>> {
let tekken_path = cache_repo.get("tekken.json");
if let Some(tekken_path) = tekken_path {
return Ok(Some(TokenizerSource::Tekken(tekken_path)));
}
let Some(tokenizer_path) = cache_repo.get("tokenizer.json").or_else(|| {
// tiktoken.model is the most common name, try it first.
cache_repo.get("tiktoken.model").or_else(|| {
// Scan for any *.tiktoken file in the cache snapshot directory.
let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf();
discover_tiktoken_in_dir(&snapshot_dir)
let Some(tokenizer_path) = cache_repo
.get("tekken.json")
.or_else(|| cache_repo.get("tokenizer.json"))
.or_else(|| {
// tiktoken.model is the most common name, try it first.
cache_repo.get("tiktoken.model").or_else(|| {
// Scan for any *.tiktoken file in the cache snapshot directory.
let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf();
discover_tiktoken_in_dir(&snapshot_dir)
})
})
}) else {
else {
return Ok(None);
};
Ok(Some(resolve_tokenizer_source(
tokenizer_path,
tokenizer_config.tokenizer_class.as_deref(),
None,
)))
)?))
}
fn resolve_local_tokenizer_source(
model_dir: &Path,
tokenizer_config: &HfTokenizerConfig,
tokenizer_config: &TokenizerConfig,
) -> Result<TokenizerSource> {
let tekken_path = local_file_if_exists(model_dir, "tekken.json");
if let Some(tekken_path) = tekken_path {
return Ok(TokenizerSource::Tekken(tekken_path));
}
let tokenizer_path = local_file_if_exists(model_dir, "tokenizer.json")
let tokenizer_path = local_file_if_exists(model_dir, "tekken.json")
.or_else(|| local_file_if_exists(model_dir, "tokenizer.json"))
.or_else(|| local_file_if_exists(model_dir, "tiktoken.model"))
.or_else(|| discover_tiktoken_in_dir(model_dir))
.ok_or_else(|| {
Error::Tokenizer(format!(
Error::new(format!(
"local model directory '{}' does not contain a supported tokenizer file \
(tokenizer.json, tiktoken.model, or *.tiktoken)",
model_dir.display()
))
})?;
Ok(resolve_tokenizer_source(
tokenizer_path,
tokenizer_config.tokenizer_class.as_deref(),
None,
))
resolve_tokenizer_source(tokenizer_path, tokenizer_config.tokenizer_class.as_deref())
}
/// Choose the tokenizer.
///
/// Selection order:
/// 1. `tekken.json`Mistral native tokenizer (preferred over HF `tokenizer.json` because the HF
/// version has a known regex bug for Mistral models).
/// 2. File extension — `.tiktoken` / `tiktoken.model` files use tiktoken from BPE data.
/// 3. `tokenizer_class` in `tokenizer_config.json` — classes containing "Tiktoken" (case-
/// 1. File extension — `.tiktoken` / `tiktoken.model` files use tiktoken from BPE data.
/// 2. `tokenizer_class` in `tokenizer_config.json` — classes containing "Tiktoken" (case-
/// insensitive) trigger tiktoken loading from a sibling BPE file.
/// 4. Default — `tokenizer.json` in HuggingFace format.
/// 3. Default — `tokenizer.json` in HuggingFace format.
fn resolve_tokenizer_source(
tokenizer_path: PathBuf,
tokenizer_class: Option<&str>,
tekken_path: Option<PathBuf>,
) -> TokenizerSource {
if let Some(tekken_path) = tekken_path {
return TokenizerSource::Tekken(tekken_path);
}
) -> Result<TokenizerSource> {
let tokenizer = TokenizerSource::from_path(tokenizer_path)?;
if is_tiktoken_file(&tokenizer_path) {
return TokenizerSource::Tiktoken(tokenizer_path);
}
if tokenizer_class.is_some_and(|cls| cls.to_ascii_lowercase().contains("tiktoken"))
&& let Some(dir) = tokenizer_path.parent()
if let TokenizerSource::HuggingFace(path) = &tokenizer
&& tokenizer_class.is_some_and(|cls| cls.to_ascii_lowercase().contains("tiktoken"))
&& let Some(dir) = path.parent()
&& let Some(tiktoken_path) = discover_tiktoken_in_dir(dir)
{
return TokenizerSource::Tiktoken(tiktoken_path);
return Ok(TokenizerSource::Tiktoken(tiktoken_path));
}
TokenizerSource::HuggingFace(tokenizer_path)
Ok(tokenizer)
}
/// Download `filename` only if it exists in `siblings`.
@@ -286,7 +289,7 @@ async fn download_if_present(
async fn download_known_file(repo: &ApiRepo, model_id: &str, filename: &str) -> Result<PathBuf> {
repo.get(filename).await.map_err(|error| {
Error::Tokenizer(format!(
Error::new(format!(
"failed to download '{filename}' for model '{model_id}': {}",
error.as_report()
))
@@ -317,7 +320,7 @@ fn find_tiktoken_sibling<'a>(siblings: &std::collections::BTreeSet<&'a str>) ->
}
/// Discover a tiktoken model file in a local directory.
pub(super) fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
let tiktoken_model = dir.join("tiktoken.model");
if tiktoken_model.exists() {
return Some(tiktoken_model);
@@ -337,7 +340,7 @@ pub(super) fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf>
}
/// Returns `true` if `path` points to a tiktoken-format file (by name).
pub(super) fn is_tiktoken_file(path: &std::path::Path) -> bool {
fn is_tiktoken_file(path: &std::path::Path) -> bool {
path.file_name()
.and_then(|n| n.to_str())
.is_some_and(|name| name == "tiktoken.model" || name.ends_with(".tiktoken"))
@@ -368,7 +371,6 @@ mod tests {
use std::fs;
use tempfile::tempdir;
use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
use super::{ResolvedModelFiles, TokenizerSource};
@@ -399,61 +401,4 @@ mod tests {
Some(dir.path().join("tokenizer_config.json"))
);
}
#[tokio::test]
#[ignore = "too slow for CI and requires network access to Hugging Face"]
async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
.await
.expect("resolve real Kimi K2.5 model files");
let tokenizer_path = match &files.tokenizer {
TokenizerSource::Tiktoken(path) => path.clone(),
other => panic!("expected tiktoken tokenizer source, got {other:?}"),
};
for backend in [
TiktokenTokenizer::new_riptoken(&tokenizer_path).expect("load riptoken backend"),
TiktokenTokenizer::new_tiktoken_rs(&tokenizer_path).expect("load tiktoken-rs backend"),
] {
let think_id = backend.token_to_id("<think>").expect("resolve <think>");
let end_think_id = backend.token_to_id("</think>").expect("resolve </think>");
let tool_section_id = backend
.token_to_id("<|tool_calls_section_begin|>")
.expect("resolve tool call section marker");
let contraction_heavy_text =
"I'm sure it's fine, but I can't say I'd trust that it's what we'd ship.";
let contraction_heavy_ids = backend.encode(contraction_heavy_text, false).unwrap();
assert_eq!(
(think_id, end_think_id, tool_section_id),
(163606, 163607, 163595)
);
assert_eq!(backend.decode(&[think_id], true).unwrap(), "<think>");
assert_eq!(backend.decode(&[end_think_id], true).unwrap(), "</think>");
assert_eq!(
backend.decode(&[tool_section_id], true).unwrap(),
"<|tool_calls_section_begin|>"
);
// This demonstrates that we're using Kimi's custom BPE pattern.
// With CL100K this will be 23 tokens instead.
assert_eq!(
contraction_heavy_ids,
vec![
17172, 3287, 4643, 8201, 11, 996, 374, 8971, 3637, 20020, 8173, 473, 4643,
1573, 56229, 13922, 13,
]
);
assert_eq!(contraction_heavy_ids.len(), 17);
assert_eq!(
backend.decode(&contraction_heavy_ids, false).unwrap(),
contraction_heavy_text
);
// Special-looking text that is not actually registered should fail gracefully.
assert_eq!(backend.token_to_id("◁think▷"), None);
assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None);
}
}
}
@@ -238,6 +238,9 @@ mod tests {
let error = parser.parse_chunk(&input).unwrap_err();
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
expect![[
r#"tool parser parsing failed: near "tool<tool▁sep>get_weather\n```json\n{}": "#
]]
.assert_eq(&error.to_report_string());
}
}
@@ -242,6 +242,7 @@ mod tests {
let error = parser.parse_chunk(&input).unwrap_err();
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
expect![[r#"tool parser parsing failed: near "<tool▁sep>{}": "#]]
.assert_eq(&error.to_report_string());
}
}
+4 -2
View File
@@ -540,8 +540,10 @@ mod tests {
.parse_chunk(r#"<tool_call>{"name":"f","arguments":42}</tool_call>"#)
.unwrap_err();
expect!["tool parser parsing failed: invalid Granite4 arguments"]
.assert_eq(&error.to_report_string());
expect![[
r#"tool parser parsing failed: near "42}</tool_call>": invalid Granite4 arguments"#
]]
.assert_eq(&error.to_report_string());
}
#[test]
+1 -1
View File
@@ -335,7 +335,7 @@ mod tests {
let error = parser.parse_chunk(&input).unwrap_err();
expect![[r#"
tool parser parsing failed: invalid InternLM2
tool parser parsing failed: near "{\"name\":\"get_weather\",\"params\":{\"location\":\"Tokyo\"}}<|action_end|>": invalid InternLM2
expected `parameters`, `arguments`"#]]
.assert_eq(&error.to_report_string());
}
+3 -3
View File
@@ -336,7 +336,7 @@ mod tests {
.unwrap_err();
expect![[r#"
tool parser parsing failed: invalid Llama JSON
tool parser parsing failed: near "{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Tokyo\"}}": invalid Llama JSON
expected `parameters`"#]]
.assert_eq(&error.to_report_string());
}
@@ -474,7 +474,7 @@ mod tests {
let error = parser.parse_chunk(r#"{"parameters":{},"name":"get_weather"}"#).unwrap_err();
expect![[r#"
tool parser parsing failed: invalid Llama JSON
tool parser parsing failed: near "{\"parameters\":{},\"name\":\"get_weather\"}": invalid Llama JSON
expected `name`"#]]
.assert_eq(&error.to_report_string());
}
@@ -489,7 +489,7 @@ mod tests {
))
.unwrap_err();
expect!["tool parser parsing failed: invalid Llama JSON"]
expect![[r#"tool parser parsing failed: near " trailing": invalid Llama JSON"#]]
.assert_eq(&error.to_report_string());
}
}
+1 -1
View File
@@ -240,7 +240,7 @@ mod tests {
.unwrap_err();
expect![[r#"
tool parser parsing failed: invalid Mistral
tool parser parsing failed: near "{\"arguments\":{},\"name\":\"get_weather\"}]": invalid Mistral
expected `name`"#]]
.assert_eq(&error.to_report_string());
}
+1 -1
View File
@@ -280,7 +280,7 @@ mod tests {
.unwrap_err();
expect![[r#"
tool parser parsing failed: invalid Qwen XML
tool parser parsing failed: near "{\"arguments\":{},\"name\":\"get_weather\"}\n</tool_call>": invalid Qwen XML
expected `name`"#]]
.assert_eq(&error.to_report_string());
}
+4 -1
View File
@@ -594,6 +594,9 @@ mod tests {
let error = parser.parse_chunk(&input).unwrap_err();
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
expect![[
r#"tool parser parsing failed: near "get_weather<|tool_call_argument_begin|>{}": "#
]]
.assert_eq(&error.to_report_string());
}
}
+2 -1
View File
@@ -593,6 +593,7 @@ mod tests {
let mut parser = MinimaxM2ToolParser::new(&test_tools());
let error = parser.parse_chunk("<minimax:tool_call><bad></minimax:tool_call>").unwrap_err();
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
expect![[r#"tool parser parsing failed: near "<bad></minimax:tool_call>": "#]]
.assert_eq(&error.to_report_string());
}
}
+4 -1
View File
@@ -878,7 +878,10 @@ mod tests {
))
.unwrap_err();
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
expect![[
r#"tool parser parsing failed: near "]<]minimax[>[<bad>]<]minimax[>[</tool_call>": "#
]]
.assert_eq(&error.to_report_string());
}
#[test]
+3 -2
View File
@@ -685,7 +685,8 @@ mod tests {
let mut parser = Qwen3CoderToolParser::new(&test_tools());
let error = parser.parse_chunk("<tool_call>\n<bad>\n</tool_call>").unwrap_err();
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
expect![[r#"tool parser parsing failed: near "\n<bad>\n</tool_call>": "#]]
.assert_eq(&error.to_report_string());
}
#[test]
@@ -697,7 +698,7 @@ mod tests {
)
.unwrap_err();
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
expect![[r#"tool parser parsing failed: near "\n<function=get_weather>\n<parameter=location>SF</function>\n</tool_call>": "#]].assert_eq(&error.to_report_string());
}
#[test]
+28 -4
View File
@@ -395,9 +395,9 @@ pub fn parse_buffered_event<E>(
Ok(event) => event,
Err(ErrMode::Incomplete(_)) => return Ok(None),
Err(ErrMode::Backtrack(e) | ErrMode::Cut(e)) => {
// TODO: enrich context for error reporting
let snippet = buffer.char_indices().nth(80).map_or(buffer, |(i, _)| &buffer[..i]);
return Err(ToolParserError::ParsingFailed {
message: e.to_string(),
message: format!("near {snippet:?}: {e}"),
});
}
};
@@ -423,8 +423,9 @@ mod tests {
use winnow::stream::{Offset, Partial, Stream};
use super::{
JsonObjectScanState, JsonStringScanState, MarkerScanState, json_str, partial_prefix_len,
safe_text_len, safe_text_len_mul, take_json_object, take_json_string, take_until_marker,
JsonObjectScanState, JsonStringScanState, MarkerScanState, json_str, parse_buffered_event,
partial_prefix_len, safe_text_len, safe_text_len_mul, take_json_object, take_json_string,
take_until_marker,
};
#[test]
@@ -832,4 +833,27 @@ mod tests {
assert!(matches!(error, ErrMode::Incomplete(_)));
}
#[test]
fn parse_buffered_event_error_includes_input_snippet() {
let result = parse_buffered_event(" {\"x\":1}", |input| {
take_json_object(input, &mut JsonObjectScanState::default())
});
let err = result.unwrap_err().to_string();
assert!(err.contains("near \""), "error must include snippet");
}
#[test]
fn parse_buffered_event_error_truncates_long_input() {
let long_input = format!(" {}", "x".repeat(100));
let result = parse_buffered_event(&long_input, |input| {
take_json_object(input, &mut JsonObjectScanState::default())
});
let err = result.unwrap_err().to_string();
assert!(err.contains("near \""), "error must include snippet");
assert!(
!err.contains(&long_input),
"snippet must be truncated for long input"
);
}
}
@@ -199,6 +199,7 @@ mod tests {
use super::prepare_completion_request;
use crate::lora::LoraModelResolution;
use crate::routes::openai::completions::types::CompletionRequest;
use crate::routes::openai::utils::types::Normalizable;
use crate::utils::{ResolvedRequestContext, resolve_request_context};
fn request_context(headers: &HeaderMap, request_id: Option<&str>) -> ResolvedRequestContext {
@@ -249,6 +250,28 @@ mod tests {
assert!(request.ignore_eos);
}
#[test]
fn normalize_coerces_null_max_tokens_to_default() {
// An absent `max_tokens` already gets the serde default.
let absent: CompletionRequest =
serde_json::from_value(base_request_json()).expect("parse request");
assert_eq!(absent.max_tokens, Some(16));
// An explicit `null` deserializes to `None`, bypassing the default;
// `normalize` must coerce it back to match Python vLLM.
let mut request: CompletionRequest = serde_json::from_value(json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"prompt": "hello",
"stream": true,
"max_tokens": null
}))
.expect("parse request");
assert_eq!(request.max_tokens, None);
request.normalize();
assert_eq!(request.max_tokens, Some(16));
}
#[test]
fn prepare_completion_request_maps_sampling_fields() {
let request: CompletionRequest = serde_json::from_value(json!({
@@ -179,7 +179,17 @@ pub struct CompletionRequest {
pub other: Map<String, Value>,
}
impl Normalizable for CompletionRequest {}
impl Normalizable for CompletionRequest {
/// Normalize the request by applying defaults.
fn normalize(&mut self) {
// An explicit `"max_tokens": null` deserializes to `None`, bypassing the
// serde field default. Coerce it back to the default so it behaves like
// an absent field, matching Python vLLM's `normalize_null_max_tokens`.
if self.max_tokens.is_none() {
self.max_tokens = default_completion_max_tokens();
}
}
}
/// Mirrors the Python vLLM `CompletionResponse` class.
#[serde_with::skip_serializing_none]
+1 -1
View File
@@ -10,7 +10,6 @@ asynk-strim-attr.workspace = true
easy-ext.workspace = true
enum-as-inner.workspace = true
futures.workspace = true
hf-hub.workspace = true
itertools.workspace = true
reqwest.workspace = true
serde.workspace = true
@@ -22,6 +21,7 @@ tracing.workspace = true
trait-set.workspace = true
vllm-engine-core-client.workspace = true
vllm-llm.workspace = true
vllm-model-files.workspace = true
vllm-tokenizer.workspace = true
[dev-dependencies]
+12 -41
View File
@@ -1,24 +1,18 @@
use std::collections::BTreeSet;
use std::fs;
use std::path::Path;
use serde::{Deserialize, Serialize};
use thiserror_ext::AsReport as _;
use vllm_model_files::read_json_file;
use crate::error::{Error, Result};
/// Minimal subset of `tokenizer_config.json` needed by chat/EOS handling.
#[derive(Debug, Default, Deserialize)]
#[serde(default)]
pub struct HfTokenizerConfig {
pub struct TokenizerConfig {
#[serde(flatten)]
pub special_tokens: HfSpecialTokens,
pub special_tokens: SpecialTokens,
pub chat_template: Option<String>,
/// The `tokenizer_class` field from HuggingFace tokenizer configs. Some
/// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value
/// containing "Tiktoken" which can be used as a hint for backend
/// selection.
pub tokenizer_class: Option<String>,
}
/// Hugging Face named special tokens may be serialized as a string or an
@@ -61,14 +55,14 @@ impl NamedSpecialToken {
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
#[serde(default)]
pub struct HfSpecialTokens {
pub struct SpecialTokens {
pub bos_token: Option<NamedSpecialToken>,
pub eos_token: Option<NamedSpecialToken>,
pub unk_token: Option<NamedSpecialToken>,
pub pad_token: Option<NamedSpecialToken>,
}
impl HfSpecialTokens {
impl SpecialTokens {
/// Returns true if we don't discover any special tokens in the config.
pub fn is_empty(&self) -> bool {
self.bos_token.is_none()
@@ -234,42 +228,19 @@ impl ModelConfig {
}
}
/// Load the tokenizer-side EOS metadata if a config file is present.
pub fn load_tokenizer_config(path: Option<&Path>) -> Result<HfTokenizerConfig> {
read_json_file(path)
}
/// Load the generation-side EOS metadata if a config file is present.
pub(super) fn load_generation_config(path: Option<&Path>) -> Result<GenerationConfig> {
read_json_file(path)
Ok(read_json_file(path)?)
}
/// Load the tokenizer-side EOS metadata if a config file is present.
pub fn load_tokenizer_config(path: Option<&Path>) -> Result<TokenizerConfig> {
Ok(read_json_file(path)?)
}
/// Load the model-side config (`config.json`) if present.
pub fn load_model_config(path: Option<&Path>) -> Result<ModelConfig> {
read_json_file(path)
}
fn read_json_file<T>(path: Option<&Path>) -> Result<T>
where
T: for<'de> Deserialize<'de> + Default,
{
let Some(path) = path else {
return Ok(T::default());
};
let content = fs::read_to_string(path).map_err(|error| {
Error::Tokenizer(format!(
"failed to read {}: {}",
path.display(),
error.as_report()
))
})?;
serde_json::from_str(&content).map_err(|error| {
Error::Tokenizer(format!(
"failed to parse {}: {}",
path.display(),
error.as_report()
))
})
Ok(read_json_file(path)?)
}
#[cfg(test)]
+66 -3
View File
@@ -1,5 +1,4 @@
mod config;
mod model_files;
use std::collections::BTreeSet;
use std::sync::Arc;
@@ -9,12 +8,12 @@ use vllm_tokenizer::{DynTokenizer, HuggingFaceTokenizer, TekkenTokenizer, Tiktok
use self::config::{GenerationConfig, load_generation_config};
pub use self::config::{
HfSpecialTokens, HfTokenizerConfig, ModelConfig, NamedSpecialToken, load_model_config,
ModelConfig, NamedSpecialToken, SpecialTokens, TokenizerConfig, load_model_config,
load_tokenizer_config,
};
pub use self::model_files::{ResolvedModelFiles, TokenizerSource};
use crate::backend::{SamplingHints, TextBackend};
use crate::error::Result;
pub use vllm_model_files::{ResolvedModelFiles, TokenizerSource};
fn load_tokenizer(tokenizer: &TokenizerSource) -> Result<DynTokenizer> {
match tokenizer {
@@ -125,3 +124,67 @@ impl TextBackend for HfTextBackend {
})
}
}
#[cfg(test)]
mod tests {
use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
use super::{ResolvedModelFiles, TokenizerSource};
#[tokio::test]
#[ignore = "too slow for CI and requires network access to Hugging Face"]
async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
.await
.expect("resolve real Kimi K2.5 model files");
let tokenizer_path = match &files.tokenizer {
TokenizerSource::Tiktoken(path) => path.clone(),
other => panic!("expected tiktoken tokenizer source, got {other:?}"),
};
for backend in [
TiktokenTokenizer::new_riptoken(&tokenizer_path).expect("load riptoken backend"),
TiktokenTokenizer::new_tiktoken_rs(&tokenizer_path).expect("load tiktoken-rs backend"),
] {
let think_id = backend.token_to_id("<think>").expect("resolve <think>");
let end_think_id = backend.token_to_id("</think>").expect("resolve </think>");
let tool_section_id = backend
.token_to_id("<|tool_calls_section_begin|>")
.expect("resolve tool call section marker");
let contraction_heavy_text =
"I'm sure it's fine, but I can't say I'd trust that it's what we'd ship.";
let contraction_heavy_ids = backend.encode(contraction_heavy_text, false).unwrap();
assert_eq!(
(think_id, end_think_id, tool_section_id),
(163606, 163607, 163595)
);
assert_eq!(backend.decode(&[think_id], true).unwrap(), "<think>");
assert_eq!(backend.decode(&[end_think_id], true).unwrap(), "</think>");
assert_eq!(
backend.decode(&[tool_section_id], true).unwrap(),
"<|tool_calls_section_begin|>"
);
// This demonstrates that we're using Kimi's custom BPE pattern.
// With CL100K this will be 23 tokens instead.
assert_eq!(
contraction_heavy_ids,
vec![
17172, 3287, 4643, 8201, 11, 996, 374, 8971, 3637, 20020, 8173, 473, 4643,
1573, 56229, 13922, 13,
]
);
assert_eq!(contraction_heavy_ids.len(), 17);
assert_eq!(
backend.decode(&contraction_heavy_ids, false).unwrap(),
contraction_heavy_text
);
// Special-looking text that is not actually registered should fail gracefully.
assert_eq!(backend.token_to_id("◁think▷"), None);
assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None);
}
}
}
+3
View File
@@ -1,6 +1,7 @@
use thiserror::Error;
use vllm_engine_core_client::Error as EngineCoreError;
use vllm_llm::Error as LlmError;
use vllm_model_files::Error as ModelFilesError;
pub use crate::lower::logprobs::LogprobsError;
pub use crate::lower::token_ids::TokenIdsError;
@@ -20,6 +21,8 @@ pub enum Error {
Logprobs(#[from] LogprobsError),
#[error(transparent)]
TokenIds(#[from] TokenIdsError),
#[error(transparent)]
ModelFiles(#[from] ModelFilesError),
#[error(
"`min_tokens` must be less than or equal to `max_tokens`, \
got min_tokens={min_tokens}, max_tokens={max_tokens}"
@@ -122,7 +122,7 @@ class QKNormRoPETestModel(torch.nn.Module):
@pytest.mark.parametrize("enable_rope_custom_op", [True])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.skipif(
not current_platform.is_cuda_alike(),
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
reason="Only test on cuda and rocm platform",
)
def test_qk_norm_rope_fusion(
@@ -136,7 +136,7 @@ def test_qk_norm_rope_fusion(
if not hasattr(torch.ops._C, "fused_qk_norm_rope"):
pytest.skip("fused_qk_norm_rope custom op not available")
torch.set_default_device("cuda")
torch.set_default_device(current_platform.device_type)
torch.set_default_dtype(dtype)
torch.manual_seed(0)
+223
View File
@@ -0,0 +1,223 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import random
import typing
import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import vllm.envs as envs
from tests.utils import ensure_current_vllm_config
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
from vllm.distributed.device_communicators.pynccl_allocator import (
get_nccl_mem_pool,
is_symmetric_memory_enabled,
)
from vllm.distributed.parallel_state import (
get_tp_group,
init_distributed_environment,
initialize_model_parallel,
)
from vllm.platforms import current_platform
from vllm.utils.system_utils import update_environment_variables
torch.manual_seed(42)
random.seed(44)
test_size_elements = 4 * 1024 * 1024
def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
monkeypatch = pytest.MonkeyPatch()
with monkeypatch.context() as m:
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
dtype = torch.bfloat16
device = torch.device(f"cuda:{local_rank}")
torch.accelerator.set_device_index(device)
torch.set_default_device(device)
torch.set_default_dtype(dtype)
update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": "12345",
}
)
init_distributed_environment()
with ensure_current_vllm_config():
initialize_model_parallel(tensor_model_parallel_size=world_size)
cuda_communicator = typing.cast(
CudaCommunicator, get_tp_group().device_communicator
)
pynccl_comm = cuda_communicator.pynccl_comm
if get_nccl_mem_pool() is None:
pytest.skip(
"NCCL allocator compilation failed (probably missing NCCL headers)."
)
if not is_symmetric_memory_enabled():
pytest.skip("NCCL symmetric memory allreduce is disabled.")
register_nccl_symmetric_ops(pynccl_comm)
input = torch.randint(1, 23, (test_size_elements,), dtype=dtype, device=device)
input_clone = input.clone()
output = torch.ops.vllm.all_reduce_symmetric_with_copy(input)
assert output is not None
group = get_tp_group().device_group
dist.all_reduce(input_clone, group=group)
torch.testing.assert_close(output, input_clone, atol=2.5, rtol=0.1)
@pytest.mark.skipif(
not current_platform.is_cuda(),
reason="NCCLSymmMemAllreduce is only available for CUDA platforms.",
)
@pytest.mark.parametrize("world_size", [2])
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
def test_nccl_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch, world_size):
if world_size > torch.accelerator.device_count():
pytest.skip("Not enough GPUs to run the test.")
# Enable SymmMemCommunicator
monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
mp.spawn(nccl_symm_mem_allreduce_worker, args=(world_size,), nprocs=world_size)
cleanup_dist_env_and_memory()
def nccl_symm_mem_allgather_worker(local_rank: int, world_size: int):
monkeypatch = pytest.MonkeyPatch()
with monkeypatch.context() as m:
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
dtype = torch.bfloat16
device = torch.device(f"cuda:{local_rank}")
torch.accelerator.set_device_index(device)
torch.set_default_device(device)
torch.set_default_dtype(dtype)
update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": "12346",
}
)
init_distributed_environment()
with ensure_current_vllm_config():
initialize_model_parallel(tensor_model_parallel_size=world_size)
cuda_communicator = typing.cast(
CudaCommunicator, get_tp_group().device_communicator
)
if get_nccl_mem_pool() is None:
pytest.skip(
"NCCL allocator compilation failed (probably missing NCCL headers)."
)
if not is_symmetric_memory_enabled():
pytest.skip("NCCL symmetric memory is disabled.")
per_rank_size = test_size_elements // world_size
input_tensor = torch.randint(
1, 23, (per_rank_size,), dtype=dtype, device=device
)
output = cuda_communicator.all_gatherv(input_tensor, dim=0)
group = get_tp_group().device_group
expected = torch.empty(test_size_elements, dtype=dtype, device=device)
dist.all_gather_into_tensor(expected, input_tensor, group=group)
torch.testing.assert_close(output, expected, atol=0.0, rtol=0.0)
@pytest.mark.skipif(
not current_platform.is_cuda(),
reason="NCCL symmetric memory is only available for CUDA platforms.",
)
@pytest.mark.parametrize("world_size", [2])
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
def test_nccl_symm_mem_allgather(monkeypatch: pytest.MonkeyPatch, world_size):
if world_size > torch.accelerator.device_count():
pytest.skip("Not enough GPUs to run the test.")
monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
mp.spawn(nccl_symm_mem_allgather_worker, args=(world_size,), nprocs=world_size)
cleanup_dist_env_and_memory()
def nccl_symm_mem_reduce_scatter_worker(local_rank: int, world_size: int):
monkeypatch = pytest.MonkeyPatch()
with monkeypatch.context() as m:
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
dtype = torch.bfloat16
device = torch.device(f"cuda:{local_rank}")
torch.accelerator.set_device_index(device)
torch.set_default_device(device)
torch.set_default_dtype(dtype)
update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": "12347",
}
)
init_distributed_environment()
with ensure_current_vllm_config():
initialize_model_parallel(tensor_model_parallel_size=world_size)
cuda_communicator = typing.cast(
CudaCommunicator, get_tp_group().device_communicator
)
if get_nccl_mem_pool() is None:
pytest.skip(
"NCCL allocator compilation failed (probably missing NCCL headers)."
)
if not is_symmetric_memory_enabled():
pytest.skip("NCCL symmetric memory is disabled.")
per_rank_size = test_size_elements // world_size
input_tensor = torch.randint(
1, 23, (test_size_elements,), dtype=dtype, device=device
)
input_clone = input_tensor.clone()
output = cuda_communicator.reduce_scatter(input_tensor, dim=0)
group = get_tp_group().device_group
expected = torch.empty(per_rank_size, dtype=dtype, device=device)
dist.reduce_scatter_tensor(expected, input_clone, group=group)
torch.testing.assert_close(output, expected, atol=2.5, rtol=0.1)
@pytest.mark.skipif(
not current_platform.is_cuda(),
reason="NCCL symmetric memory is only available for CUDA platforms.",
)
@pytest.mark.parametrize("world_size", [2])
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
def test_nccl_symm_mem_reduce_scatter(monkeypatch: pytest.MonkeyPatch, world_size):
if world_size > torch.accelerator.device_count():
pytest.skip("Not enough GPUs to run the test.")
monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
mp.spawn(nccl_symm_mem_reduce_scatter_worker, args=(world_size,), nprocs=world_size)
cleanup_dist_env_and_memory()
@@ -1,96 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import random
import typing
import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import vllm.envs as envs
from tests.utils import ensure_current_vllm_config
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
from vllm.distributed.device_communicators.pynccl_allocator import (
get_nccl_mem_pool,
is_symmetric_memory_enabled,
)
from vllm.distributed.parallel_state import (
get_tp_group,
init_distributed_environment,
initialize_model_parallel,
)
from vllm.platforms import current_platform
from vllm.utils.system_utils import update_environment_variables
torch.manual_seed(42)
random.seed(44)
test_size_elements = 4 * 1024 * 1024
def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
monkeypatch = pytest.MonkeyPatch()
with monkeypatch.context() as m:
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
dtype = torch.bfloat16
device = torch.device(f"cuda:{local_rank}")
torch.accelerator.set_device_index(device)
torch.set_default_device(device)
torch.set_default_dtype(dtype)
update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": "12345",
}
)
init_distributed_environment()
with ensure_current_vllm_config():
initialize_model_parallel(tensor_model_parallel_size=world_size)
cuda_communicator = typing.cast(
CudaCommunicator, get_tp_group().device_communicator
)
pynccl_comm = cuda_communicator.pynccl_comm
if get_nccl_mem_pool() is None:
pytest.skip(
"NCCL allocator compilation failed (probably missing NCCL headers)."
)
if not is_symmetric_memory_enabled():
pytest.skip("NCCL symmetric memory allreduce is disabled.")
register_nccl_symmetric_ops(pynccl_comm)
input = torch.randint(1, 23, (test_size_elements,), dtype=dtype, device=device)
input_clone = input.clone()
output = torch.ops.vllm.all_reduce_symmetric_with_copy(input)
assert output is not None
group = get_tp_group().device_group
dist.all_reduce(input_clone, group=group)
torch.testing.assert_close(output, input_clone, atol=2.5, rtol=0.1)
@pytest.mark.skipif(
not current_platform.is_cuda(),
reason="NCCLSymmMemAllreduce is only available for CUDA platforms.",
)
@pytest.mark.parametrize("world_size", [2])
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
def test_nccl_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch, world_size):
if world_size > torch.accelerator.device_count():
pytest.skip("Not enough GPUs to run the test.")
# Enable SymmMemCommunicator
monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
mp.spawn(nccl_symm_mem_allreduce_worker, args=(world_size,), nprocs=world_size)
cleanup_dist_env_and_memory()
@@ -115,7 +115,6 @@ TEXT_GENERATION_MODELS = {
"google/gemma-1.1-2b-it": PPTestSettings.fast(),
"google/gemma-2-9b": PPTestSettings.fast(),
"gpt2": PPTestSettings.fast(),
"bigcode/starcoder": PPTestSettings.fast(),
"EleutherAI/gpt-j-6b": PPTestSettings.fast(),
"EleutherAI/pythia-1.4b": PPTestSettings.fast(),
"ibm/PowerLM-3b": PPTestSettings.fast(),
File diff suppressed because it is too large Load Diff
@@ -251,7 +251,13 @@ def _validate_field_consistency(events: list) -> None:
"response.reasoning_part.added",
):
_assert_item_fields(event, etype, active_item_id, active_output_index)
active_content_index = getattr(event, "content_index", None)
content_index = getattr(event, "content_index", None)
if active_content_index is None:
assert content_index == 0, (
f"{etype} for a new item must start at content_index 0, "
f"got {content_index}"
)
active_content_index = content_index
continue
# --- all other item-level events --------------------------
@@ -454,6 +454,7 @@ async def test_streaming(client: OpenAI, model_name: str, background: bool):
if event.type == "response.output_item.added":
assert event.item.id != current_item_id
current_item_id = event.item.id
current_content_index = -1
elif event.type in [
"response.output_text.delta",
"response.reasoning_text.delta",
@@ -465,7 +466,7 @@ async def test_streaming(client: OpenAI, model_name: str, background: bool):
"response.content_part.added",
"response.reasoning_part.added",
]:
assert event.content_index != current_content_index
assert event.content_index == current_content_index + 1
current_content_index = event.content_index
elif event.type in [
"response.output_text.delta",
@@ -2,8 +2,10 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for vllm.entrypoints.openai.responses.harmony."""
import pytest
from openai.types.responses import (
ResponseFunctionToolCall,
ResponseFunctionWebSearch,
ResponseOutputMessage,
ResponseReasoningItem,
)
@@ -12,7 +14,6 @@ from openai_harmony import Author, Message, Role, TextContent
from vllm.entrypoints.openai.responses.harmony import (
harmony_to_response_output,
parser_state_to_response_output,
response_previous_input_to_harmony,
)
@@ -95,7 +96,8 @@ class TestResponsePreviousInputToHarmony:
class TestHarmonyToResponseOutput:
"""Tests for harmony_to_response_output function."""
def test_commentary_with_no_recipient_creates_message(self):
@pytest.mark.parametrize("incomplete", [False, True])
def test_commentary_with_no_recipient_creates_message(self, incomplete):
"""Test that commentary with recipient=None (preambles) creates message items.
Per Harmony format, preambles are intended to be shown to end-users,
@@ -108,13 +110,15 @@ class TestHarmonyToResponseOutput:
message = message.with_channel("commentary")
# recipient is None by default, representing a preamble
output_items = harmony_to_response_output(message)
output_items = harmony_to_response_output(
message, frozenset(), incomplete=incomplete
)
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseOutputMessage)
assert output_items[0].type == "message"
assert output_items[0].role == "assistant"
assert output_items[0].status == "completed"
assert output_items[0].status == ("incomplete" if incomplete else "completed")
assert len(output_items[0].content) == 1
assert output_items[0].content[0].type == "output_text"
assert (
@@ -122,82 +126,148 @@ class TestHarmonyToResponseOutput:
== "I will now search for the weather information."
)
def test_commentary_with_function_recipient_creates_function_call(self):
"""Test commentary with recipient='functions.X' creates function calls."""
message = Message.from_role_and_content(
Role.ASSISTANT, '{"location": "San Francisco", "units": "celsius"}'
)
message = message.with_channel("commentary")
message = message.with_recipient("functions.get_weather")
@pytest.mark.parametrize("channel", ["commentary", "comment", "analysis", "final"])
@pytest.mark.parametrize(
("recipient", "fn_names", "expected_name"),
[
("functions.get_weather", frozenset(), "get_weather"),
("get_weather", frozenset({"get_weather"}), "get_weather"),
("math.sum", frozenset({"math.sum"}), "math.sum"),
],
)
@pytest.mark.parametrize("incomplete", [False, True])
def test_function_recipient_creates_function_call(
self, channel, recipient, fn_names, expected_name, incomplete
):
"""Function recipients create function calls across channels."""
content = '{"location": "San Francisco"}'
if recipient == "math.sum":
content = '{"a": 1, "b": 2}'
output_items = harmony_to_response_output(message)
message = Message.from_role_and_content(Role.ASSISTANT, content)
message = message.with_channel(channel)
message = message.with_recipient(recipient)
output_items = harmony_to_response_output(
message, fn_names, incomplete=incomplete
)
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseFunctionToolCall)
assert output_items[0].type == "function_call"
assert output_items[0].name == "get_weather"
assert (
output_items[0].arguments
== '{"location": "San Francisco", "units": "celsius"}'
)
assert output_items[0].name == expected_name
assert output_items[0].arguments == content
assert output_items[0].call_id.startswith("call_")
assert output_items[0].id.startswith("fc_")
assert output_items[0].status == ("incomplete" if incomplete else "completed")
@pytest.mark.parametrize("channel", ["commentary", "comment", "analysis", "final"])
@pytest.mark.parametrize(
("recipient", "content"),
[
("python", "import numpy as np\nprint(np.array([1, 2, 3]))"),
("browser", "Navigating to the specified URL"),
("container", "Running command in container"),
],
)
@pytest.mark.parametrize("incomplete", [False, True])
def test_builtin_recipient_creates_reasoning(
self, channel, recipient, content, incomplete
):
"""Built-in recipients create reasoning items."""
message = Message.from_role_and_content(Role.ASSISTANT, content)
message = message.with_channel(channel)
message = message.with_recipient(recipient)
output_items = harmony_to_response_output(
message, frozenset(), incomplete=incomplete
)
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseReasoningItem)
assert output_items[0].type == "reasoning"
assert output_items[0].content[0].text == content
assert output_items[0].status is None
@pytest.mark.parametrize("channel", ["commentary", "comment", "analysis", "final"])
@pytest.mark.parametrize(
("recipient", "fn_names", "content", "expected_name", "expected_server_label"),
[
(
"get_weather",
frozenset(),
'{"arg": "value"}',
"get_weather",
"get_weather",
),
(
"not_get_weather",
frozenset({"get_weather"}),
'{"arg": "value"}',
"not_get_weather",
"not_get_weather",
),
("repo_browser.list", frozenset(), '{"cmd": "ls"}', "list", "repo_browser"),
],
)
@pytest.mark.parametrize("incomplete", [False, True])
def test_non_function_non_builtin_recipient_creates_mcp_call(
self,
channel,
recipient,
fn_names,
content,
expected_name,
expected_server_label,
incomplete,
):
"""Non-function, non-built-in recipients create MCP calls."""
message = Message.from_role_and_content(Role.ASSISTANT, content)
message = message.with_channel(channel)
message = message.with_recipient(recipient)
output_items = harmony_to_response_output(
message, fn_names, incomplete=incomplete
)
assert len(output_items) == 1
assert isinstance(output_items[0], McpCall)
assert output_items[0].type == "mcp_call"
assert output_items[0].name == expected_name
assert output_items[0].server_label == expected_server_label
assert output_items[0].arguments == content
assert output_items[0].status == ("incomplete" if incomplete else "completed")
@pytest.mark.parametrize("incomplete", [False, True])
def test_browser_search_recipient_respects_incomplete(self, incomplete):
"""browser.search emits a web search call unless the item is incomplete."""
message = Message.from_role_and_content(
Role.ASSISTANT, '{"query": "weather in San Francisco"}'
)
message = message.with_channel("commentary")
message = message.with_recipient("browser.search")
output_items = harmony_to_response_output(
message, frozenset(), incomplete=incomplete
)
if incomplete:
assert output_items == []
return
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseFunctionWebSearch)
assert output_items[0].type == "web_search_call"
assert output_items[0].status == "completed"
def test_commentary_with_python_recipient_creates_reasoning(self):
"""Test that commentary with recipient='python' creates reasoning items."""
message = Message.from_role_and_content(
Role.ASSISTANT, "import numpy as np\nprint(np.array([1, 2, 3]))"
)
message = message.with_channel("commentary")
message = message.with_recipient("python")
output_items = harmony_to_response_output(message)
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseReasoningItem)
assert output_items[0].type == "reasoning"
assert (
output_items[0].content[0].text
== "import numpy as np\nprint(np.array([1, 2, 3]))"
)
def test_commentary_with_browser_recipient_creates_reasoning(self):
"""Test that commentary with recipient='browser' creates reasoning items."""
message = Message.from_role_and_content(
Role.ASSISTANT, "Navigating to the specified URL"
)
message = message.with_channel("commentary")
message = message.with_recipient("browser")
output_items = harmony_to_response_output(message)
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseReasoningItem)
assert output_items[0].type == "reasoning"
assert output_items[0].content[0].text == "Navigating to the specified URL"
def test_commentary_with_container_recipient_creates_reasoning(self):
"""Test that commentary with recipient='container' creates reasoning items."""
message = Message.from_role_and_content(
Role.ASSISTANT, "Running command in container"
)
message = message.with_channel("commentary")
message = message.with_recipient("container")
output_items = harmony_to_response_output(message)
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseReasoningItem)
assert output_items[0].type == "reasoning"
assert output_items[0].content[0].text == "Running command in container"
assert output_items[0].action.type == "search"
assert output_items[0].action.query == "cursor:weather in San Francisco"
def test_commentary_with_empty_content_and_no_recipient(self):
"""Test edge case: empty commentary with recipient=None."""
message = Message.from_role_and_content(Role.ASSISTANT, "")
message = message.with_channel("commentary")
output_items = harmony_to_response_output(message)
output_items = harmony_to_response_output(message, frozenset())
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseOutputMessage)
@@ -212,7 +282,7 @@ class TestHarmonyToResponseOutput:
message = Message.from_role_and_contents(Role.ASSISTANT, contents)
message = message.with_channel("commentary")
output_items = harmony_to_response_output(message)
output_items = harmony_to_response_output(message, frozenset())
# _parse_final_message returns single ResponseOutputMessage with
# multiple contents
@@ -232,7 +302,7 @@ class TestHarmonyToResponseOutput:
message = message.with_channel("commentary")
message = message.with_recipient("functions.get_weather")
output_items = harmony_to_response_output(message)
output_items = harmony_to_response_output(message, frozenset())
assert len(output_items) == 2
assert all(isinstance(item, ResponseFunctionToolCall) for item in output_items)
@@ -241,21 +311,6 @@ class TestHarmonyToResponseOutput:
assert output_items[0].arguments == '{"location": "San Francisco"}'
assert output_items[1].arguments == '{"location": "New York"}'
def test_commentary_with_unknown_recipient_creates_mcp_call(self):
"""Test that commentary with unknown recipient creates MCP call."""
message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
message = message.with_channel("commentary")
message = message.with_recipient("custom_tool")
fn_names = frozenset({"other_tool"})
output_items = harmony_to_response_output(message, fn_names)
assert len(output_items) == 1
assert isinstance(output_items[0], McpCall)
assert output_items[0].type == "mcp_call"
assert output_items[0].name == "custom_tool"
assert output_items[0].server_label == "custom_tool"
def test_analysis_channel_creates_reasoning(self):
"""Test that analysis channel creates reasoning items."""
message = Message.from_role_and_content(
@@ -263,7 +318,7 @@ class TestHarmonyToResponseOutput:
)
message = message.with_channel("analysis")
output_items = harmony_to_response_output(message)
output_items = harmony_to_response_output(message, frozenset())
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseReasoningItem)
@@ -283,352 +338,6 @@ class TestHarmonyToResponseOutput:
"The weather is sunny, 72°F",
)
output_items = harmony_to_response_output(message)
output_items = harmony_to_response_output(message, frozenset())
assert len(output_items) == 0
class TestHarmonyToResponseOutputWithFunctionToolNames:
"""Tests for bare function name handling with function_tool_names."""
def test_bare_name_creates_function_call_when_in_tool_names(self):
"""Bare function name matching a known tool creates function call."""
message = Message.from_role_and_content(
Role.ASSISTANT, '{"location": "San Francisco"}'
)
message = message.with_channel("commentary")
message = message.with_recipient("get_weather")
fn_names = frozenset({"get_weather"})
output_items = harmony_to_response_output(message, fn_names)
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseFunctionToolCall)
assert output_items[0].type == "function_call"
assert output_items[0].name == "get_weather"
assert output_items[0].arguments == '{"location": "San Francisco"}'
def test_bare_name_creates_mcp_call_when_not_in_tool_names(self):
"""Bare name not matching any known tool creates MCP call."""
message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
message = message.with_channel("commentary")
message = message.with_recipient("custom_tool")
fn_names = frozenset({"get_weather"})
output_items = harmony_to_response_output(message, fn_names)
assert len(output_items) == 1
assert isinstance(output_items[0], McpCall)
assert output_items[0].type == "mcp_call"
def test_dotted_function_name_creates_function_call(self):
"""Dotted function name in tool names creates function call."""
message = Message.from_role_and_content(Role.ASSISTANT, '{"a": 1, "b": 2}')
message = message.with_channel("commentary")
message = message.with_recipient("math.sum")
fn_names = frozenset({"math.sum"})
output_items = harmony_to_response_output(message, fn_names)
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseFunctionToolCall)
assert output_items[0].name == "math.sum"
def test_empty_tool_names_defaults_to_mcp(self):
"""With empty function_tool_names, bare names become MCP calls."""
message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
message = message.with_channel("commentary")
message = message.with_recipient("get_weather")
output_items = harmony_to_response_output(message, frozenset())
assert len(output_items) == 1
assert isinstance(output_items[0], McpCall)
def test_prefixed_name_always_function_call(self):
"""functions. prefix always creates function call even with empty tool names."""
message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
message = message.with_channel("commentary")
message = message.with_recipient("functions.get_weather")
output_items = harmony_to_response_output(message, frozenset())
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseFunctionToolCall)
assert output_items[0].name == "get_weather"
class TestParserStateWithFunctionToolNames:
"""Tests for parser_state_to_response_output with function_tool_names."""
def test_bare_name_creates_function_call(self):
from unittest.mock import Mock
parser = Mock()
parser.current_content = '{"arg": "value"}'
parser.current_role = Role.ASSISTANT
parser.current_channel = "commentary"
parser.current_recipient = "get_weather"
fn_names = frozenset({"get_weather"})
items = parser_state_to_response_output(parser, fn_names)
assert len(items) == 1
assert isinstance(items[0], ResponseFunctionToolCall)
assert items[0].name == "get_weather"
assert items[0].status == "in_progress"
def test_bare_name_creates_mcp_when_not_in_tool_names(self):
from unittest.mock import Mock
parser = Mock()
parser.current_content = '{"arg": "value"}'
parser.current_role = Role.ASSISTANT
parser.current_channel = "commentary"
parser.current_recipient = "unknown_tool"
fn_names = frozenset({"get_weather"})
items = parser_state_to_response_output(parser, fn_names)
assert len(items) == 1
assert isinstance(items[0], McpCall)
assert items[0].name == "unknown_tool"
class TestToolCallsOnNonStandardChannels:
"""Tests verifying tool calls are detected regardless of channel."""
def test_function_call_on_comment_channel(self):
message = Message.from_role_and_content(Role.ASSISTANT, '{"query": "weather"}')
message = message.with_channel("comment")
message = message.with_recipient("functions.get_weather")
output_items = harmony_to_response_output(message)
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseFunctionToolCall)
assert output_items[0].type == "function_call"
assert output_items[0].name == "get_weather"
def test_bare_function_on_comment_channel(self):
message = Message.from_role_and_content(Role.ASSISTANT, '{"query": "weather"}')
message = message.with_channel("comment")
message = message.with_recipient("get_weather")
fn_names = frozenset({"get_weather"})
output_items = harmony_to_response_output(message, fn_names)
assert len(output_items) == 1
assert isinstance(output_items[0], ResponseFunctionToolCall)
assert output_items[0].name == "get_weather"
def test_parser_state_comment_channel_function(self):
from unittest.mock import Mock
parser = Mock()
parser.current_content = '{"arg": "value"}'
parser.current_role = Role.ASSISTANT
parser.current_channel = "comment"
parser.current_recipient = "functions.get_weather"
items = parser_state_to_response_output(parser)
assert len(items) == 1
assert isinstance(items[0], ResponseFunctionToolCall)
assert items[0].name == "get_weather"
def test_parser_state_comment_channel_mcp(self):
from unittest.mock import Mock
parser = Mock()
parser.current_content = '{"arg": "value"}'
parser.current_role = Role.ASSISTANT
parser.current_channel = "comment"
parser.current_recipient = "mcp.server.tool"
fn_names: frozenset[str] = frozenset()
items = parser_state_to_response_output(parser, fn_names)
assert len(items) == 1
assert isinstance(items[0], McpCall)
def test_parse_mcp_call_basic() -> None:
"""Test that MCP calls are parsed with correct type and server_label."""
message = Message.from_role_and_content(Role.ASSISTANT, '{"path": "/tmp"}')
message = message.with_recipient("filesystem")
message = message.with_channel("commentary")
fn_names: frozenset[str] = frozenset()
output_items = harmony_to_response_output(message, fn_names)
assert len(output_items) == 1
assert isinstance(output_items[0], McpCall)
assert output_items[0].type == "mcp_call"
assert output_items[0].name == "filesystem"
assert output_items[0].server_label == "filesystem"
assert output_items[0].arguments == '{"path": "/tmp"}'
assert output_items[0].status == "completed"
def test_parse_mcp_call_dotted_recipient() -> None:
"""Test that dotted recipients extract the tool name correctly."""
message = Message.from_role_and_content(Role.ASSISTANT, '{"cmd": "ls"}')
message = message.with_recipient("repo_browser.list")
message = message.with_channel("commentary")
fn_names: frozenset[str] = frozenset()
output_items = harmony_to_response_output(message, fn_names)
assert len(output_items) == 1
assert isinstance(output_items[0], McpCall)
assert output_items[0].name == "list"
assert output_items[0].server_label == "repo_browser"
def test_mcp_vs_function_call() -> None:
"""Test that function calls are not parsed as MCP calls."""
func_message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
func_message = func_message.with_recipient("functions.my_tool")
func_message = func_message.with_channel("commentary")
func_items = harmony_to_response_output(func_message)
assert len(func_items) == 1
assert not isinstance(func_items[0], McpCall)
assert func_items[0].type == "function_call"
def test_mcp_vs_builtin_tools() -> None:
"""Test that built-in tools (python, container) are not parsed as MCP calls."""
# Test python (built-in tool) - should be reasoning, not MCP
python_message = Message.from_role_and_content(Role.ASSISTANT, "print('hello')")
python_message = python_message.with_recipient("python")
python_message = python_message.with_channel("commentary")
python_items = harmony_to_response_output(python_message)
assert len(python_items) == 1
assert not isinstance(python_items[0], McpCall)
assert python_items[0].type == "reasoning"
def test_parser_state_to_response_output_commentary_channel() -> None:
"""Test parser_state_to_response_output with commentary
channel and various recipients."""
from unittest.mock import Mock
# Test 1: functions.* recipient -> should return function tool call
parser_func = Mock()
parser_func.current_content = '{"arg": "value"}'
parser_func.current_role = Role.ASSISTANT
parser_func.current_channel = "commentary"
parser_func.current_recipient = "functions.my_tool"
func_items = parser_state_to_response_output(parser_func)
assert len(func_items) == 1
assert not isinstance(func_items[0], McpCall)
assert func_items[0].type == "function_call"
assert func_items[0].name == "my_tool"
assert func_items[0].status == "in_progress"
# Test 2: MCP tool (not builtin) -> should return MCP call
parser_mcp = Mock()
parser_mcp.current_content = '{"path": "/tmp"}'
parser_mcp.current_role = Role.ASSISTANT
parser_mcp.current_channel = "commentary"
parser_mcp.current_recipient = "filesystem"
fn_names: frozenset[str] = frozenset()
mcp_items = parser_state_to_response_output(parser_mcp, fn_names)
assert len(mcp_items) == 1
assert isinstance(mcp_items[0], McpCall)
assert mcp_items[0].type == "mcp_call"
assert mcp_items[0].name == "filesystem"
assert mcp_items[0].server_label == "filesystem"
assert mcp_items[0].status == "in_progress"
# Test 3: Built-in tool (python)
# should NOT return MCP call, returns reasoning (internal tool interaction)
parser_builtin = Mock()
parser_builtin.current_content = "print('hello')"
parser_builtin.current_role = Role.ASSISTANT
parser_builtin.current_channel = "commentary"
parser_builtin.current_recipient = "python"
builtin_items = parser_state_to_response_output(parser_builtin)
# Built-in tools explicitly return reasoning
assert len(builtin_items) == 1
assert not isinstance(builtin_items[0], McpCall)
assert builtin_items[0].type == "reasoning"
# Test 4: No recipient (preamble) → should return message, not reasoning
parser_preamble = Mock()
parser_preamble.current_content = "I'll search for that information now."
parser_preamble.current_role = Role.ASSISTANT
parser_preamble.current_channel = "commentary"
parser_preamble.current_recipient = None
preamble_items = parser_state_to_response_output(parser_preamble)
assert len(preamble_items) == 1
assert isinstance(preamble_items[0], ResponseOutputMessage)
assert preamble_items[0].type == "message"
assert preamble_items[0].content[0].text == "I'll search for that information now."
assert preamble_items[0].status == "incomplete" # streaming
def test_parser_state_to_response_output_analysis_channel() -> None:
"""Test parser_state_to_response_output with analysis
channel and various recipients."""
from unittest.mock import Mock
# Test 1: functions.* recipient -> should return function tool call
parser_func = Mock()
parser_func.current_content = '{"arg": "value"}'
parser_func.current_role = Role.ASSISTANT
parser_func.current_channel = "analysis"
parser_func.current_recipient = "functions.my_tool"
func_items = parser_state_to_response_output(parser_func)
assert len(func_items) == 1
assert not isinstance(func_items[0], McpCall)
assert func_items[0].type == "function_call"
assert func_items[0].name == "my_tool"
assert func_items[0].status == "in_progress"
# Test 2: MCP tool (not builtin) -> should return MCP call
parser_mcp = Mock()
parser_mcp.current_content = '{"query": "test"}'
parser_mcp.current_role = Role.ASSISTANT
parser_mcp.current_channel = "analysis"
parser_mcp.current_recipient = "database"
fn_names: frozenset[str] = frozenset()
mcp_items = parser_state_to_response_output(parser_mcp, fn_names)
assert len(mcp_items) == 1
assert isinstance(mcp_items[0], McpCall)
assert mcp_items[0].type == "mcp_call"
assert mcp_items[0].name == "database"
assert mcp_items[0].server_label == "database"
assert mcp_items[0].status == "in_progress"
# Test 3: Built-in tool (container)
# should NOT return MCP call, falls through to reasoning
parser_builtin = Mock()
parser_builtin.current_content = "docker run"
parser_builtin.current_role = Role.ASSISTANT
parser_builtin.current_channel = "analysis"
parser_builtin.current_recipient = "container"
builtin_items = parser_state_to_response_output(parser_builtin)
# Should fall through to reasoning logic
assert len(builtin_items) == 1
assert not isinstance(builtin_items[0], McpCall)
assert builtin_items[0].type == "reasoning"
@@ -51,6 +51,7 @@ from vllm.entrypoints.openai.responses.streaming_events import (
)
from vllm.inputs import tokens_input
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.parser.harmony import Segment
from vllm.sampling_params import SamplingParams
@@ -534,13 +535,9 @@ class TestHarmonyPreambleStreaming:
"""Tests for preamble (commentary with no recipient) streaming events."""
@staticmethod
def _make_ctx(*, channel, recipient, delta="hello"):
"""Build a lightweight mock StreamingHarmonyContext."""
ctx = MagicMock()
ctx.last_content_delta = delta
ctx.parser.current_channel = channel
ctx.parser.current_recipient = recipient
return ctx
def _make_segment(*, channel, recipient, delta="hello"):
"""Build a lightweight segment for Harmony streaming tests."""
return Segment(channel=channel, recipient=recipient, delta=delta)
@staticmethod
def _make_previous_item(*, channel, recipient, text="preamble text"):
@@ -559,10 +556,10 @@ class TestHarmonyPreambleStreaming:
emit_content_delta_events,
)
ctx = self._make_ctx(channel="commentary", recipient=None)
segment = self._make_segment(channel="commentary", recipient=None)
state = StreamingState()
events = emit_content_delta_events(ctx, state)
events = emit_content_delta_events(segment, state)
type_names = [e.type for e in events]
assert "response.output_text.delta" in type_names
@@ -574,13 +571,13 @@ class TestHarmonyPreambleStreaming:
emit_content_delta_events,
)
ctx = self._make_ctx(channel="commentary", recipient=None, delta="w")
segment = self._make_segment(channel="commentary", recipient=None, delta="w")
state = StreamingState()
state.sent_output_item_added = True
state.current_item_id = "msg_test"
state.current_content_index = 0
events = emit_content_delta_events(ctx, state)
events = emit_content_delta_events(segment, state)
type_names = [e.type for e in events]
assert "response.output_text.delta" in type_names
@@ -592,13 +589,13 @@ class TestHarmonyPreambleStreaming:
emit_content_delta_events,
)
ctx = self._make_ctx(
segment = self._make_segment(
channel="commentary",
recipient="functions.get_weather",
)
state = StreamingState()
events = emit_content_delta_events(ctx, state)
events = emit_content_delta_events(segment, state)
type_names = [e.type for e in events]
assert "response.output_text.delta" not in type_names
@@ -612,6 +609,7 @@ class TestHarmonyPreambleStreaming:
previous = self._make_previous_item(channel="commentary", recipient=None)
state = StreamingState()
state.sent_output_item_added = True
state.current_item_id = "msg_test"
state.current_output_index = 0
state.current_content_index = 0
@@ -634,13 +632,53 @@ class TestHarmonyPreambleStreaming:
channel="commentary", recipient="functions.get_weather"
)
state = StreamingState()
state.is_first_function_call_delta = True
state.current_item_id = "fc_test"
state.current_call_id = "call_test"
events = emit_previous_item_done_events(previous, state)
type_names = [e.type for e in events]
assert "response.output_text.done" not in type_names
@pytest.mark.xfail(
reason=(
"TODO: Ensure added/in-progress events are emitted for zero-delta items."
"So we can safely emit done events for zero-delta items."
),
strict=True,
)
def test_zero_delta_items_should_preserve_streaming_lifecycle(
self,
) -> None:
"""Zero-delta Harmony items should still produce a coherent lifecycle."""
from vllm.entrypoints.openai.responses.streaming_events import (
emit_previous_item_done_events,
)
cases: list[tuple[str, str | None, str]] = [
("commentary", None, "msg_stale"),
("analysis", None, "msg_stale"),
("commentary", "functions.get_weather", "fc_stale"),
("commentary", "python", "tool_stale"),
("commentary", "repo_browser.list", "mcp_stale"),
]
for channel, recipient, current_item_id in cases:
previous = self._make_previous_item(channel=channel, recipient=recipient)
state = StreamingState()
state.current_item_id = current_item_id
state.current_call_id = "call_stale"
state.current_content_index = 0
events = emit_previous_item_done_events(
previous, state, function_tool_names=None
)
type_names = [e.type for e in events]
assert "response.output_item.added" in type_names
assert "response.output_item.done" in type_names
def _make_simple_context_with_output(text, token_ids, response_parser=None):
"""Create a SimpleContext with a RequestOutput containing the given text."""
@@ -176,6 +176,49 @@ def test_build_multi_port_external_lb_child_args_sets_external_rank_server():
assert child_args.api_server_count == 1
def test_run_vllm_dp_server_uses_python_server_by_default(monkeypatch):
calls: list[str] = []
monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None)
monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None)
monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None)
monkeypatch.setattr(dp_sup.envs, "VLLM_RUST_FRONTEND_PATH", None, raising=False)
monkeypatch.setattr(
dp_sup, "_run_python_vllm_dp_server", lambda _args: calls.append("python")
)
monkeypatch.setattr(
dp_sup, "_run_rust_vllm_dp_server", lambda _args: calls.append("rust")
)
dp_sup._run_vllm_dp_server(_make_unit_args(data_parallel_rank=4))
assert calls == ["python"]
def test_run_vllm_dp_server_uses_rust_frontend_when_enabled(monkeypatch):
calls: list[str] = []
monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None)
monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None)
monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None)
monkeypatch.setattr(
dp_sup.envs,
"VLLM_RUST_FRONTEND_PATH",
"/tmp/vllm-rs",
raising=False,
)
monkeypatch.setattr(
dp_sup, "_run_python_vllm_dp_server", lambda _args: calls.append("python")
)
monkeypatch.setattr(
dp_sup, "_run_rust_vllm_dp_server", lambda _args: calls.append("rust")
)
dp_sup._run_vllm_dp_server(_make_unit_args(data_parallel_rank=4))
assert calls == ["rust"]
def test_validate_multi_port_external_lb_args_allows_ssl():
args = _make_unit_args(
ssl_keyfile="/tmp/server.key",
@@ -12,15 +12,21 @@ Tests real customer usage scenarios:
- Priority: env vars > decorators > customer script files > framework
defaults
Note: These tests focus on validating server responses rather than directly calling
get_ping_handler() and get_invoke_handler() to ensure full integration testing.
The handler-override scenarios exercise the real vLLM SageMaker router and
bootstrap path via an in-process FastAPI ``TestClient`` instead of launching a
model server. These scenarios fully replace the ``/ping`` and ``/invocations``
endpoints with customer handlers, so no inference engine is required to
validate override behavior. Avoiding the model server also keeps the tests
fast and deterministic rather than depending on the FastAPI version resolved
into the test environment at runtime.
"""
import os
import tempfile
import pytest
import requests
from fastapi import FastAPI
from fastapi.testclient import TestClient
from tests.utils import RemoteOpenAIServer
@@ -29,6 +35,23 @@ from .conftest import (
)
def _build_sagemaker_test_client() -> TestClient:
"""Build a TestClient over the real SageMaker router and bootstrap path.
``attach_router`` is called with empty supported tasks because the override
tests replace the endpoints with customer handlers, so no framework
invocation handler (and therefore no engine) is exercised.
"""
from vllm.entrypoints.serve.sagemaker.api_router import (
attach_router,
sagemaker_standards_bootstrap,
)
app = FastAPI()
attach_router(app, ())
return TestClient(sagemaker_standards_bootstrap(app))
class TestHandlerOverrideIntegration:
"""Integration tests simulating real customer usage scenarios.
@@ -89,8 +112,7 @@ class TestHandlerOverrideIntegration:
except ImportError:
pass
@pytest.mark.asyncio
async def test_customer_script_functions_auto_loaded(self):
def test_customer_script_functions_auto_loaded(self, monkeypatch, tmp_path):
"""Test customer scenario: script functions automatically override
framework defaults."""
try:
@@ -101,15 +123,15 @@ class TestHandlerOverrideIntegration:
pytest.skip("model-hosting-container-standards not available")
# Customer writes a script file with ping() and invoke() functions
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
f.write(
"""
script_path = tmp_path / "model.py"
script_path.write_text(
"""
from fastapi import Request
async def custom_sagemaker_ping_handler():
return {
"status": "healthy",
"source": "customer_override",
"source": "customer_override",
"message": "Custom ping from customer script"
}
@@ -119,62 +141,39 @@ async def custom_sagemaker_invocation_handler(request: Request):
"source": "customer_override"
}
"""
)
# Customer sets SageMaker environment variables to point to their script
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
with _build_sagemaker_test_client() as client:
# Customer tests their server and sees their overrides work
# automatically
ping_response = client.get("/ping")
assert ping_response.status_code == 200
ping_data = ping_response.json()
invoke_response = client.post(
"/invocations",
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
script_path = f.name
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
try:
script_dir = os.path.dirname(script_path)
script_name = os.path.basename(script_path)
# Customer sets SageMaker environment variables to point to their script
env_vars = {
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
}
args = [
"--dtype",
"bfloat16",
"--max-model-len",
"2048",
"--enforce-eager",
"--max-num-seqs",
"32",
# Customer sees their functions are used
assert ping_data["source"] == "customer_override"
assert ping_data["message"] == "Custom ping from customer script"
assert invoke_data["source"] == "customer_override"
assert invoke_data["predictions"] == [
"Custom response from customer script"
]
with RemoteOpenAIServer(
MODEL_NAME_SMOLLM, args, env_dict=env_vars
) as server:
# Customer tests their server and sees their overrides work
# automatically
ping_response = requests.get(server.url_for("ping"))
assert ping_response.status_code == 200
ping_data = ping_response.json()
invoke_response = requests.post(
server.url_for("invocations"),
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
# Customer sees their functions are used
assert ping_data["source"] == "customer_override"
assert ping_data["message"] == "Custom ping from customer script"
assert invoke_data["source"] == "customer_override"
assert invoke_data["predictions"] == [
"Custom response from customer script"
]
finally:
os.unlink(script_path)
@pytest.mark.asyncio
async def test_customer_decorator_usage(self):
def test_customer_decorator_usage(self, monkeypatch, tmp_path):
"""Test customer scenario: using @custom_ping_handler and
@custom_invocation_handler decorators."""
try:
@@ -185,9 +184,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
pytest.skip("model-hosting-container-standards not available")
# Customer writes a script file with decorators
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
f.write(
"""
script_path = tmp_path / "model.py"
script_path.write_text(
"""
import model_hosting_container_standards.sagemaker as sagemaker_standards
from fastapi import Request
@@ -198,62 +197,39 @@ async def my_ping():
"source": "customer_decorator"
}
@sagemaker_standards.custom_invocation_handler
@sagemaker_standards.custom_invocation_handler
async def my_invoke(request: Request):
return {
"type": "invoke",
"type": "invoke",
"source": "customer_decorator"
}
"""
)
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
with _build_sagemaker_test_client() as client:
ping_response = client.get("/ping")
assert ping_response.status_code == 200
ping_data = ping_response.json()
invoke_response = client.post(
"/invocations",
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
script_path = f.name
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
try:
script_dir = os.path.dirname(script_path)
script_name = os.path.basename(script_path)
# Customer sees their handlers are used by the server
assert ping_data["source"] == "customer_decorator"
assert invoke_data["source"] == "customer_decorator"
env_vars = {
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
}
args = [
"--dtype",
"bfloat16",
"--max-model-len",
"2048",
"--enforce-eager",
"--max-num-seqs",
"32",
]
with RemoteOpenAIServer(
MODEL_NAME_SMOLLM, args, env_dict=env_vars
) as server:
ping_response = requests.get(server.url_for("ping"))
assert ping_response.status_code == 200
ping_data = ping_response.json()
invoke_response = requests.post(
server.url_for("invocations"),
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
# Customer sees their handlers are used by the server
assert ping_data["source"] == "customer_decorator"
assert invoke_data["source"] == "customer_decorator"
finally:
os.unlink(script_path)
@pytest.mark.asyncio
async def test_handler_priority_order(self):
def test_handler_priority_order(self, monkeypatch, tmp_path):
"""Test priority: @custom_ping_handler/@custom_invocation_handler
decorators vs script functions."""
try:
@@ -264,9 +240,9 @@ async def my_invoke(request: Request):
pytest.skip("model-hosting-container-standards not available")
# Customer writes a script with both decorator and regular functions
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
f.write(
"""
script_path = tmp_path / "model.py"
script_path.write_text(
"""
import model_hosting_container_standards.sagemaker as sagemaker_standards
from fastapi import Request
@@ -275,7 +251,7 @@ from fastapi import Request
async def decorated_ping():
return {
"status": "healthy",
"source": "ping_decorator_in_script",
"source": "ping_decorator_in_script",
"priority": "decorator"
}
@@ -296,60 +272,37 @@ async def custom_sagemaker_invocation_handler(request: Request):
"priority": "function"
}
"""
)
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
with _build_sagemaker_test_client() as client:
ping_response = client.get("/ping")
assert ping_response.status_code == 200
ping_data = ping_response.json()
invoke_response = client.post(
"/invocations",
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
script_path = f.name
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
try:
script_dir = os.path.dirname(script_path)
script_name = os.path.basename(script_path)
# @custom_ping_handler decorator has higher priority than
# script function
assert ping_data["source"] == "ping_decorator_in_script"
assert ping_data["priority"] == "decorator"
env_vars = {
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
}
# Script function is used for invoke
assert invoke_data["source"] == "script_invoke_function"
assert invoke_data["priority"] == "function"
args = [
"--dtype",
"bfloat16",
"--max-model-len",
"2048",
"--enforce-eager",
"--max-num-seqs",
"32",
]
with RemoteOpenAIServer(
MODEL_NAME_SMOLLM, args, env_dict=env_vars
) as server:
ping_response = requests.get(server.url_for("ping"))
assert ping_response.status_code == 200
ping_data = ping_response.json()
invoke_response = requests.post(
server.url_for("invocations"),
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
# @custom_ping_handler decorator has higher priority than
# script function
assert ping_data["source"] == "ping_decorator_in_script"
assert ping_data["priority"] == "decorator"
# Script function is used for invoke
assert invoke_data["source"] == "script_invoke_function"
assert invoke_data["priority"] == "function"
finally:
os.unlink(script_path)
@pytest.mark.asyncio
async def test_environment_variable_script_loading(self):
def test_environment_variable_script_loading(self, monkeypatch, tmp_path):
"""Test that environment variables correctly specify script location
and loading."""
try:
@@ -360,9 +313,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
pytest.skip("model-hosting-container-standards not available")
# Customer writes a script in a specific directory
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
f.write(
"""
script_path = tmp_path / "model.py"
script_path.write_text(
"""
from fastapi import Request
async def custom_sagemaker_ping_handler():
@@ -379,60 +332,43 @@ async def custom_sagemaker_invocation_handler(request: Request):
"method": "environment_variable_loading"
}
"""
)
# Test environment variable script loading
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
with _build_sagemaker_test_client() as client:
ping_response = client.get("/ping")
assert ping_response.status_code == 200
ping_data = ping_response.json()
invoke_response = client.post(
"/invocations",
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
script_path = f.name
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
try:
script_dir = os.path.dirname(script_path)
script_name = os.path.basename(script_path)
# Test environment variable script loading
env_vars = {
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
}
args = [
"--dtype",
"bfloat16",
"--max-model-len",
"2048",
"--enforce-eager",
"--max-num-seqs",
"32",
]
with RemoteOpenAIServer(
MODEL_NAME_SMOLLM, args, env_dict=env_vars
) as server:
ping_response = requests.get(server.url_for("ping"))
assert ping_response.status_code == 200
ping_data = ping_response.json()
invoke_response = requests.post(
server.url_for("invocations"),
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
# Verify that the script was loaded via environment variables
assert ping_data["source"] == "env_loaded_script"
assert ping_data["method"] == "environment_variable_loading"
assert invoke_data["source"] == "env_loaded_script"
assert invoke_data["method"] == "environment_variable_loading"
finally:
os.unlink(script_path)
# Verify that the script was loaded via environment variables
assert ping_data["source"] == "env_loaded_script"
assert ping_data["method"] == "environment_variable_loading"
assert invoke_data["source"] == "env_loaded_script"
assert invoke_data["method"] == "environment_variable_loading"
@pytest.mark.asyncio
async def test_framework_default_handlers(self):
"""Test that framework default handlers work when no customer
overrides exist."""
overrides exist.
This scenario exercises the real inference path (default
``/invocations``), so it keeps using a live model server rather than
the in-process TestClient.
"""
args = [
"--dtype",
"bfloat16",
@@ -478,8 +414,7 @@ async def custom_sagemaker_invocation_handler(request: Request):
)
assert invoke_response.status_code == 200
@pytest.mark.asyncio
async def test_handler_env_var_override(self):
def test_handler_env_var_override(self, monkeypatch, tmp_path):
"""Test CUSTOM_FASTAPI_PING_HANDLER and CUSTOM_FASTAPI_INVOCATION_HANDLER
environment variable overrides."""
try:
@@ -493,9 +428,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
pytest.skip("model-hosting-container-standards not available")
# Create a script with both env var handlers and script functions
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
f.write(
"""
script_path = tmp_path / "model.py"
script_path.write_text(
"""
from fastapi import Request, Response
import json
@@ -533,68 +468,47 @@ async def custom_sagemaker_invocation_handler(request: Request):
"method": "script_function"
}
"""
)
# Set environment variables to override both handlers
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
monkeypatch.setenv(
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER,
f"{script_path.name}:env_var_ping_handler",
)
monkeypatch.setenv(
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER,
f"{script_path.name}:env_var_invoke_handler",
)
with _build_sagemaker_test_client() as client:
# Test ping handler override
ping_response = client.get("/ping")
assert ping_response.status_code == 200
ping_data = ping_response.json()
# Environment variable should override script function
assert ping_data["method"] == "environment_variable"
assert ping_data["source"] == "env_var_ping"
# Test invocation handler override
invoke_response = client.post(
"/invocations",
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
script_path = f.name
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
try:
script_dir = os.path.dirname(script_path)
script_name = os.path.basename(script_path)
# Environment variable should override script function
assert invoke_data["method"] == "environment_variable"
assert invoke_data["source"] == "env_var_invoke"
# Set environment variables to override both handlers
env_vars = {
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER: (
f"{script_name}:env_var_ping_handler"
),
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER: (
f"{script_name}:env_var_invoke_handler"
),
}
args = [
"--dtype",
"bfloat16",
"--max-model-len",
"2048",
"--enforce-eager",
"--max-num-seqs",
"32",
]
with RemoteOpenAIServer(
MODEL_NAME_SMOLLM, args, env_dict=env_vars
) as server:
# Test ping handler override
ping_response = requests.get(server.url_for("ping"))
assert ping_response.status_code == 200
ping_data = ping_response.json()
# Environment variable should override script function
assert ping_data["method"] == "environment_variable"
assert ping_data["source"] == "env_var_ping"
# Test invocation handler override
invoke_response = requests.post(
server.url_for("invocations"),
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
# Environment variable should override script function
assert invoke_data["method"] == "environment_variable"
assert invoke_data["source"] == "env_var_invoke"
finally:
os.unlink(script_path)
@pytest.mark.asyncio
async def test_env_var_priority_over_decorator_and_script(self):
def test_env_var_priority_over_decorator_and_script(self, monkeypatch, tmp_path):
"""Test that environment variables have highest priority over decorators
and script functions for both ping and invocation handlers."""
try:
@@ -608,9 +522,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
pytest.skip("model-hosting-container-standards not available")
# Create a script with all three handler types for both ping and invocation
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
f.write(
"""
script_path = tmp_path / "model.py"
script_path.write_text(
"""
import model_hosting_container_standards.sagemaker as sagemaker_standards
from fastapi import Request, Response
import json
@@ -674,62 +588,42 @@ async def custom_sagemaker_invocation_handler(request: Request):
"priority": "script_function"
}
"""
)
# Set environment variables to specify highest priority handlers
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
monkeypatch.setenv(
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER,
f"{script_path.name}:env_priority_ping",
)
monkeypatch.setenv(
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER,
f"{script_path.name}:env_priority_invoke",
)
with _build_sagemaker_test_client() as client:
# Test ping handler priority
ping_response = client.get("/ping")
assert ping_response.status_code == 200
ping_data = ping_response.json()
# Environment variable has highest priority and should be used
assert ping_data["priority"] == "environment_variable"
assert ping_data["source"] == "env_var"
# Test invocation handler priority
invoke_response = client.post(
"/invocations",
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
script_path = f.name
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
try:
script_dir = os.path.dirname(script_path)
script_name = os.path.basename(script_path)
# Set environment variables to specify highest priority handlers
env_vars = {
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER: (
f"{script_name}:env_priority_ping"
),
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER: (
f"{script_name}:env_priority_invoke"
),
}
args = [
"--dtype",
"bfloat16",
"--max-model-len",
"2048",
"--enforce-eager",
"--max-num-seqs",
"32",
]
with RemoteOpenAIServer(
MODEL_NAME_SMOLLM, args, env_dict=env_vars
) as server:
# Test ping handler priority
ping_response = requests.get(server.url_for("ping"))
assert ping_response.status_code == 200
ping_data = ping_response.json()
# Environment variable has highest priority and should be used
assert ping_data["priority"] == "environment_variable"
assert ping_data["source"] == "env_var"
# Test invocation handler priority
invoke_response = requests.post(
server.url_for("invocations"),
json={
"model": MODEL_NAME_SMOLLM,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 5,
},
)
assert invoke_response.status_code == 200
invoke_data = invoke_response.json()
# Environment variable has highest priority and should be used
assert invoke_data["priority"] == "environment_variable"
assert invoke_data["source"] == "env_var"
finally:
os.unlink(script_path)
# Environment variable has highest priority and should be used
assert invoke_data["priority"] == "environment_variable"
assert invoke_data["source"] == "env_var"
+116 -89
View File
@@ -1,18 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import pytest
from openai_harmony import Author, Message, Role, StreamState, TextContent
from openai_harmony import Author, Message, Role, TextContent
from vllm.entrypoints.openai.responses.context import (
HarmonyContext,
SimpleContext,
StreamingHarmonyContext,
TurnMetrics,
)
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.parser.harmony import ChunkResult, HarmonyParser, Segment
def create_mock_request_output(
@@ -68,25 +68,59 @@ async def generate_mock_outputs(
)
@pytest.fixture
def mock_parser():
"""Set up a mock parser for tests."""
with patch(
"vllm.entrypoints.openai.responses.context.get_streamable_parser_for_assistant"
) as mock_parser_factory:
# Create a mock parser object
parser = MagicMock()
parser.messages = []
parser.current_channel = None
parser.state = StreamState.EXPECT_START
mock_parser_factory.return_value = parser
yield parser
class FakeHarmonyParser(HarmonyParser):
def __init__(self):
# Skip HarmonyParser initialization and script outputs directly.
self.reasoning_parser = None
self.tool_parser = None
self._chunk_results: list[ChunkResult] = []
self._flush_results: list[Segment | None] = []
self.processed_chunks: list[list[int]] = []
def enqueue_chunk_result(
self,
segments: list[Segment] | None = None,
reasoning_token_count: int = 0,
) -> None:
self._chunk_results.append(
ChunkResult(
segments=[] if segments is None else segments,
reasoning_token_count=reasoning_token_count,
)
)
def enqueue_flush_result(self, segment: Segment | None) -> None:
self._flush_results.append(segment)
def process_chunk(self, token_ids) -> ChunkResult:
self.processed_chunks.append(list(token_ids))
if self._chunk_results:
return self._chunk_results.pop(0)
return ChunkResult(segments=[], reasoning_token_count=0)
def flush(self) -> Segment | None:
if self._flush_results:
return self._flush_results.pop(0)
return None
def make_harmony_context(
messages=None, available_tools=None, function_tool_names=None
) -> tuple[HarmonyContext, FakeHarmonyParser]:
fake_parser = FakeHarmonyParser()
context = HarmonyContext(
messages=[] if messages is None else messages,
available_tools=[] if available_tools is None else available_tools,
function_tool_names=function_tool_names,
response_parser=fake_parser,
)
return context, fake_parser
def test_single_turn_token_counting():
"""Test token counting behavior for a single turn."""
# Create a context
context = HarmonyContext(messages=[], available_tools=[])
context, _ = make_harmony_context()
# Create a mock RequestOutput with specific token counts
mock_output = create_mock_request_output(
@@ -118,7 +152,7 @@ def test_single_turn_token_counting():
async def test_multi_turn_token_counting():
"""Test token counting behavior across multiple turns with tool output."""
# Create a context
context = HarmonyContext(messages=[], available_tools=["browser"])
context, _ = make_harmony_context(available_tools=["browser"])
# Simulate a conversation with 3 turns
# Turn 1: prefill 5, decode 3, tool 7
@@ -177,7 +211,7 @@ async def test_multi_turn_token_counting():
def test_empty_output_tokens():
"""Test behavior when RequestOutput has empty output tokens."""
context = HarmonyContext(messages=[], available_tools=[])
context, _ = make_harmony_context()
# Create a RequestOutput with empty output tokens
mock_output = create_mock_request_output(
@@ -197,7 +231,7 @@ def test_empty_output_tokens():
def test_missing_prompt_token_ids():
"""Test behavior when RequestOutput has None prompt_token_ids."""
context = HarmonyContext(messages=[], available_tools=[])
context, _ = make_harmony_context()
mock_output = create_mock_request_output(
prompt_token_ids=None, # No prompt token IDs
@@ -216,12 +250,10 @@ def test_missing_prompt_token_ids():
assert context.num_tool_output_tokens == 0
def test_reasoning_tokens_counting(mock_parser):
def test_reasoning_tokens_counting():
"""Test that reasoning tokens are counted correctly."""
context = HarmonyContext(messages=[], available_tools=[])
# Mock parser to simulate reasoning channel
mock_parser.current_channel = "analysis" # Reasoning channel
context, parser = make_harmony_context()
parser.enqueue_chunk_result(reasoning_token_count=4)
mock_output = create_mock_request_output(
prompt_token_ids=[1, 2, 3],
@@ -236,13 +268,11 @@ def test_reasoning_tokens_counting(mock_parser):
assert context.num_output_tokens == 4
def test_preamble_tokens_not_counted_as_reasoning(mock_parser):
def test_preamble_tokens_not_counted_as_reasoning():
"""Preambles (commentary with no recipient) are visible user text,
not hidden reasoning. They must NOT inflate num_reasoning_tokens."""
context = HarmonyContext(messages=[], available_tools=[])
mock_parser.current_channel = "commentary"
mock_parser.current_recipient = None # preamble
context, parser = make_harmony_context()
parser.enqueue_chunk_result(reasoning_token_count=0)
mock_output = create_mock_request_output(
prompt_token_ids=[1, 2, 3],
@@ -255,13 +285,11 @@ def test_preamble_tokens_not_counted_as_reasoning(mock_parser):
assert context.num_output_tokens == 3
def test_commentary_with_recipient_counted_as_reasoning(mock_parser):
def test_commentary_with_recipient_counted_as_reasoning():
"""Commentary directed at a tool (recipient != None) is hidden from
the user, so it should still count as reasoning tokens."""
context = HarmonyContext(messages=[], available_tools=[])
mock_parser.current_channel = "commentary"
mock_parser.current_recipient = "python"
context, parser = make_harmony_context()
parser.enqueue_chunk_result(reasoning_token_count=3)
mock_output = create_mock_request_output(
prompt_token_ids=[1, 2, 3],
@@ -276,7 +304,7 @@ def test_commentary_with_recipient_counted_as_reasoning(mock_parser):
def test_zero_tokens_edge_case():
"""Test behavior with all zero token counts."""
context = HarmonyContext(messages=[], available_tools=[])
context, _ = make_harmony_context()
# Create a request with empty lists (not None) for both prompt and
# output tokens
@@ -299,10 +327,7 @@ def test_zero_tokens_edge_case():
@pytest.mark.asyncio
async def test_single_turn_no_tool_output():
"""Test that first turn never generates tool output tokens."""
context = HarmonyContext(
messages=[],
available_tools=["browser"], # Tools available
)
context, _ = make_harmony_context(available_tools=["browser"])
# Even with large prompt in first turn, no tool tokens should be counted
mock_output = create_mock_request_output(
@@ -324,7 +349,7 @@ async def test_negative_tool_tokens_edge_case():
tokens. We should log an error and clamp the value to 0."""
# Use patch to check if logger.error was called
with patch("vllm.entrypoints.openai.responses.context.logger.error") as mock_log:
context = HarmonyContext(messages=[], available_tools=["browser"])
context, _ = make_harmony_context(available_tools=["browser"])
# First turn
mock_output1 = create_mock_request_output(
@@ -360,15 +385,15 @@ async def test_negative_tool_tokens_edge_case():
@pytest.mark.asyncio
async def test_streaming_multi_turn_token_counting(mock_parser):
async def test_streaming_multi_turn_token_counting():
"""Test token counting for streaming multi-turn conversations.
This test focuses on how StreamingHarmonyContext counts tokens in a
This test focuses on how HarmonyContext counts tokens in a
multi-turn conversation with streaming (token-by-token) outputs and
message boundaries.
"""
# Create a streaming context
context = StreamingHarmonyContext(messages=[], available_tools=["browser"])
context, parser = make_harmony_context(available_tools=["browser"])
num_prompt_tokens = [3, 8, 13]
num_output_tokens = [3, 3, 2]
@@ -413,10 +438,8 @@ async def test_streaming_multi_turn_token_counting(mock_parser):
assert context.num_tool_output_tokens == 0 # No tool output in first turn
assert context.first_tok_of_message is True # Ready for next message
# Second turn: reasoning tokens in analysis channel
mock_parser.current_channel = "analysis" # Set to reasoning channel
# First token of second turn
parser.enqueue_chunk_result(reasoning_token_count=1)
context.append_output(
create_mock_request_output(
prompt_token_ids=[
@@ -436,6 +459,7 @@ async def test_streaming_multi_turn_token_counting(mock_parser):
)
# More tokens in reasoning channel
parser.enqueue_chunk_result(reasoning_token_count=1)
context.append_output(
create_mock_request_output(
output_token_ids=[202],
@@ -443,6 +467,7 @@ async def test_streaming_multi_turn_token_counting(mock_parser):
)
)
parser.enqueue_chunk_result(reasoning_token_count=1)
context.append_output(
create_mock_request_output(
output_token_ids=[203],
@@ -460,9 +485,6 @@ async def test_streaming_multi_turn_token_counting(mock_parser):
expected_tool_tokens = 8 - 3 - 3 # = 2
assert context.num_tool_output_tokens == expected_tool_tokens
# Third turn: regular output channel
mock_parser.current_channel = "final" # Switch back to regular channel
# Third turn (with more cached tokens)
context.append_output(
create_mock_request_output(
@@ -520,13 +542,8 @@ async def test_streaming_multi_turn_token_counting(mock_parser):
@pytest.mark.asyncio
async def test_streaming_message_synchronization(mock_parser):
"""Test message synchronization logic from lines 413-417 in context.py.
This test verifies that when parser.messages contains more messages than
the context's _messages (minus initial messages), the context properly
extends its message list with the new parser messages.
"""
async def test_streaming_message_synchronization():
"""Completed messages from append-local and flush segments sync into context."""
# Create a streaming context with some initial messages
initial_messages = [
@@ -536,23 +553,30 @@ async def test_streaming_message_synchronization(mock_parser):
recipient=Role.ASSISTANT,
)
]
context = StreamingHarmonyContext(messages=initial_messages, available_tools=[])
context, parser = make_harmony_context(messages=initial_messages)
# Verify initial state
assert len(context._messages) == 1
assert context.num_init_messages == 1
# Mock parser to have more messages than context
# Simulate parser having processed 3 new messages
mock_parser.messages = [
Message(
author=Author(role=Role.ASSISTANT, name="assistant"),
content=[TextContent(text="Response 1")],
recipient=Role.USER,
),
]
response_text = "First response"
message = Message(
author=Author(role=Role.ASSISTANT, name="assistant"),
content=[TextContent(text=response_text)],
recipient=Role.USER,
)
parser.enqueue_chunk_result(
segments=[
Segment(
channel="commentary",
recipient=None,
delta="",
completed_message=message,
)
]
)
# This should trigger the message synchronization logic
# This should sync the completed message from the latest append
context.append_output(
create_mock_request_output(
prompt_token_ids=[1, 2, 3], output_token_ids=[101], finished=False
@@ -563,36 +587,39 @@ async def test_streaming_message_synchronization(mock_parser):
assert len(context._messages) == 2
# Verify the new messages were added correctly
assert context._messages[1].content[0].text == "Response 1"
assert context._messages[1].content[0].text == response_text
# Test the specific condition from line 413-414:
# len(self._messages) - self.num_init_messages < len(self.parser.messages)
messages_minus_init = len(context._messages) - context.num_init_messages
parser_messages_count = len(mock_parser.messages)
assert messages_minus_init == 1
# After synchronization, they should be equal (no longer less than)
assert messages_minus_init == parser_messages_count
response_text = "Second response"
message = Message(
author=Author(role=Role.ASSISTANT, name="assistant"),
content=[TextContent(text=response_text)],
recipient=Role.USER,
)
flush_segment = Segment(
channel="commentary",
recipient=None,
delta="",
completed_message=message,
)
parser.enqueue_flush_result(flush_segment)
# Test edge case: add one more parser message
mock_parser.messages.append(
Message(
author=Author(role=Role.ASSISTANT, name="assistant"),
content=[TextContent(text="Response 4")],
recipient=Role.USER,
# Create another output to trigger synchronization via flush()
context.append_output(
create_mock_request_output(
prompt_token_ids=[1, 2, 3], output_token_ids=[102], finished=True
)
)
# Create another output to trigger synchronization again
mock_output2 = create_mock_request_output(
prompt_token_ids=[1, 2, 3], output_token_ids=[102], finished=True
)
context.append_output(mock_output2)
# Verify the fourth message was added, num_init_messages is still 1
# Verify the flushed response was added, num_init_messages is still 1
assert len(context._messages) == 3
assert context.num_init_messages == 1
assert context._messages[2].content[0].text == "Response 4"
assert context._messages[2].content[0].text == response_text
assert context.last_append_flush_status is True
assert len(context.last_append_segments) == 1
assert context.last_append_segments[0].completed_message is message
def test_turn_metrics_copy_and_reset():
@@ -8,7 +8,6 @@ actual NCCL communication.
"""
import os
from collections.abc import Callable
from dataclasses import dataclass
from unittest.mock import patch
@@ -48,7 +47,6 @@ class MockUpdateInfo(WeightTransferUpdateInfo):
names: list[str] | None = None
dtype_names: list[str] | None = None
shapes: list[list[int]] | None = None
num_updates_list: list[int] | None = None
class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]):
@@ -59,16 +57,20 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
# Class-level tracking for verification across processes
init_transfer_engine_called: bool = False
start_called: bool = False
receive_weights_called: bool = False
finish_called: bool = False
shutdown_called: bool = False
last_init_info: MockInitInfo | None = None
last_update_info: MockUpdateInfo | None = None
def __init__(self, config, parallel_config, model):
super().__init__(config, parallel_config, model)
def __init__(self, config, vllm_config, device, model):
super().__init__(config, vllm_config, device, model)
# Reset tracking on init
MockWeightTransferEngine.init_transfer_engine_called = False
MockWeightTransferEngine.start_called = False
MockWeightTransferEngine.receive_weights_called = False
MockWeightTransferEngine.finish_called = False
MockWeightTransferEngine.shutdown_called = False
MockWeightTransferEngine.last_init_info = None
MockWeightTransferEngine.last_update_info = None
@@ -77,37 +79,28 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
MockWeightTransferEngine.init_transfer_engine_called = True
MockWeightTransferEngine.last_init_info = init_info
def receive_weights(
self,
update_info: MockUpdateInfo,
load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
) -> None:
MockWeightTransferEngine.receive_weights_called = True
MockWeightTransferEngine.last_update_info = update_info
# Simulate loading weights by calling load_weights with empty list
# (In real implementation, this would receive and load actual weights)
load_weights([])
def start_weight_update(self) -> None:
MockWeightTransferEngine.start_called = True
def receive_sparse_weights(
self,
update_info: MockUpdateInfo,
apply_patches: Callable[[list], None],
) -> None:
def finish_weight_update(self) -> None:
MockWeightTransferEngine.finish_called = True
def receive_weights(self, update_info: MockUpdateInfo) -> None:
MockWeightTransferEngine.receive_weights_called = True
MockWeightTransferEngine.last_update_info = update_info
apply_patches([])
def shutdown(self) -> None:
MockWeightTransferEngine.shutdown_called = True
def trainer_send_weights(self, *args, **kwargs):
@staticmethod
def trainer_send_weights(*args, **kwargs):
"""Mock method to simulate trainer sending weights."""
pass
def mock_create_engine(config, parallel_config, model):
def mock_create_engine(config, vllm_config, device, model):
"""Mock factory function that returns our mock engine."""
return MockWeightTransferEngine(config, parallel_config, model)
return MockWeightTransferEngine(config, vllm_config, device, model)
# --- Tests ---
@@ -208,7 +201,7 @@ def test_update_weights_calls_engine():
llm.init_weight_transfer_engine(
WeightTransferInitRequest(init_info={"test_param": "init"})
)
llm.start_weight_update(is_checkpoint_format=True)
llm.start_weight_update()
# Call update_weights
test_names = ["layer.weight", "layer.bias"]
@@ -243,61 +236,6 @@ def test_update_weights_calls_engine():
llm.finish_weight_update()
@create_new_process_for_each_test()
def test_update_weights_passes_sparse_metadata():
"""Test sparse update metadata is forwarded unchanged to the engine."""
if torch.accelerator.device_count() < 1:
pytest.skip("Need at least 1 GPU for this test")
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
with patch(
"vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
mock_create_engine,
):
llm = LLM(
model=MODEL_NAME,
enforce_eager=True,
load_format="dummy",
tensor_parallel_size=1,
weight_transfer_config=WeightTransferConfig(backend="nccl"),
)
llm.init_weight_transfer_engine(
WeightTransferInitRequest(init_info={"test_param": "init"})
)
llm.start_weight_update(is_checkpoint_format=False)
llm.update_weights(
WeightTransferUpdateRequest(
update_info={
"names": ["layer.weight"],
"dtype_names": ["bfloat16"],
"shapes": [[100]],
"num_updates_list": [3],
"update_kind": "sparse_flat",
}
)
)
def check_sparse_update_called(self):
engine = self.weight_transfer_engine
if not engine.receive_weights_called:
return None
info = engine.last_update_info
return (
info.update_kind,
info.num_updates_list,
)
results = llm.collective_rpc(check_sparse_update_called)
for result in results:
assert result == ("sparse_flat", [3])
llm.finish_weight_update()
@create_new_process_for_each_test()
def test_full_weight_transfer_flow():
"""Test the complete weight transfer flow: init -> start -> update -> finish."""
@@ -327,7 +265,7 @@ def test_full_weight_transfer_flow():
)
# Step 2: Start weight update
llm.start_weight_update(is_checkpoint_format=True)
llm.start_weight_update()
# Step 3: Update weights
llm.update_weights(
@@ -3,6 +3,7 @@ accuracy_threshold: 0.84
num_questions: 1319
num_fewshot: 5
startup_max_wait_seconds: 1200
use_chat_completions: true
server_args: >-
--enforce-eager
--max-model-len 4096
+59 -9
View File
@@ -110,6 +110,39 @@ async def call_vllm_api(
return "", 0
async def call_vllm_chat_api(
session: aiohttp.ClientSession,
model: str,
prompt: str,
temperature: float,
max_tokens: int,
stop: list[str] | None = None,
url: str | None = None,
seed: int | None = None,
) -> tuple[str, int]:
"""Call vLLM's OpenAI-compatible chat completions endpoint."""
data = {
"model": model,
"messages": [{"role": "user", "content": prompt}],
"temperature": temperature,
"max_tokens": max_tokens,
"stop": stop,
}
if seed is not None:
data["seed"] = seed
try:
async with session.post(f"{url}/v1/chat/completions", json=data) as response:
response.raise_for_status()
result = await response.json()
text = result["choices"][0]["message"]["content"] or ""
completion_tokens = result.get("usage", {}).get("completion_tokens", 0)
return text, completion_tokens
except Exception as e:
print(f"Error calling vLLM chat API ({type(e).__name__}): {e}")
return "", 0
def _build_gsm8k_prompts(
num_questions: int = 1319,
num_shots: int = 5,
@@ -173,6 +206,8 @@ def evaluate_gsm8k(
num_questions: int = 1319,
num_shots: int = 5,
max_tokens: int = 256,
model: str | None = None,
use_chat_completions: bool = False,
host: str = "http://127.0.0.1",
port: int = 8000,
temperature: float = 0.0,
@@ -193,15 +228,30 @@ def evaluate_gsm8k(
output_tokens: list[int] = [0] * num_questions
async def get_answer(session: aiohttp.ClientSession, i: int) -> tuple[str, int]:
answer, tokens = await call_vllm_api(
session=session,
prompt=prompts[i],
temperature=temperature,
max_tokens=max_tokens,
stop=["Question", "Assistant:", "<|separator|>"],
url=base_url,
seed=seed,
)
stop = ["Question", "Assistant:", "<|separator|>"]
if use_chat_completions:
if model is None:
raise ValueError("model is required for chat completions")
answer, tokens = await call_vllm_chat_api(
session=session,
model=model,
prompt=prompts[i],
temperature=temperature,
max_tokens=max_tokens,
stop=stop,
url=base_url,
seed=seed,
)
else:
answer, tokens = await call_vllm_api(
session=session,
prompt=prompts[i],
temperature=temperature,
max_tokens=max_tokens,
stop=stop,
url=base_url,
seed=seed,
)
states[i] = answer
output_tokens[i] = tokens
return answer, tokens
@@ -65,6 +65,8 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict:
num_questions=eval_config["num_questions"],
num_shots=eval_config["num_fewshot"],
max_tokens=eval_config.get("max_tokens", 256),
model=eval_config["model_name"],
use_chat_completions=eval_config.get("use_chat_completions", False),
host=host,
port=port,
request_timeout_seconds=request_timeout_seconds,
+101
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
@@ -109,3 +110,103 @@ def test_cpu_unary_activation(
if not (activation_cls is GELU and dtype != torch.bfloat16):
raw_out = torch.empty_like(x)
opcheck(fn, (raw_out, x, *op_args))
@pytest.mark.parametrize("dtype", DTYPES)
@torch.inference_mode()
def test_cpu_gelu_tanh_and_mul(
default_vllm_config,
dtype: torch.dtype,
) -> None:
gate = torch.tensor(
[
[
-12.0,
-10.0,
-9.01,
-5.0,
-2.0,
-1.0,
-0.0,
0.0,
0.5,
1.0,
2.0,
5.0,
9.01,
10.0,
12.0,
11.0,
],
[
-7.5,
-4.5,
-3.0,
-1.5,
-0.75,
-0.25,
0.25,
0.75,
1.5,
3.0,
4.5,
7.5,
-11.0,
11.0,
8.75,
-8.75,
],
],
dtype=dtype,
)
val = torch.tensor(
[
[
0.25,
-0.5,
0.75,
-1.0,
1.25,
-1.5,
1.75,
-2.0,
2.25,
-2.5,
2.75,
-3.0,
3.25,
-3.5,
3.75,
-4.0,
],
[
-0.4,
0.6,
-0.8,
1.0,
-1.2,
1.4,
-1.6,
1.8,
-2.0,
2.2,
-2.4,
2.6,
-2.8,
3.0,
-3.2,
3.4,
],
],
dtype=dtype,
)
x = torch.cat((val, gate), dim=-1).contiguous()
kernel_out = torch.empty_like(val)
torch.ops._C.gelu_tanh_and_mul(kernel_out, x)
torch_ref = torch.nn.functional.gelu(val, approximate="tanh") * gate
atol = get_default_atol(kernel_out)
rtol = get_default_rtol(kernel_out)
torch.testing.assert_close(kernel_out, torch_ref, atol=atol, rtol=rtol)
@@ -292,6 +292,76 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
assert (topk == 7).all(), "topk buffer should be untouched on shared layer"
@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
def test_fused_norm_rope_ds_mla(num_tokens: int):
"""fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100).
Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, dynamic float32
scale) | 4 float32 scales | 64 bf16 (unquantized) RoPE.
"""
torch.manual_seed(5)
dev = "cuda"
max_pos = 8192
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
bs = max_pos
mla_cache = torch.zeros(1, bs, 656, device=dev, dtype=torch.uint8)
slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
q_out = K.fused_norm_rope(
pos,
q_c,
qw,
EPS,
kv_c,
kvw,
EPS,
k_pe,
mla_cos_sin,
None,
None,
None,
EPS,
None,
topk,
slot_mapping=slot,
indexer_k_cache=None,
mla_kv_cache=mla_cache,
mla_kv_cache_dtype="fp8_ds_mla",
mla_k_scale=None,
has_indexer=False,
index_rope_interleave=False,
)
assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (ds_mla)")
kv_ref = rms_norm(kv_c, kvw) # [N, 512] fp32
kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) # [N, 64]
tiles = kv_ref.view(num_tokens, 4, 128)
ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1.1754944e-38)
ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, KV_LORA).to(FP8)
cache = mla_cache[0, :num_tokens] # [N, 656] uint8
nope = cache[:, :KV_LORA].view(FP8)
scales = cache.view(torch.float32)[:, KV_LORA // 4 : KV_LORA // 4 + 4]
rope_off = KV_LORA // 2 + 8
rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + ROPE_DIM]
torch.testing.assert_close(scales, ref_scale, rtol=1e-2, atol=1e-6)
assert_fp8(nope, ref_nope, "ds_mla NoPE fp8")
assert_bf16(rope_vals, kpe_ref, "ds_mla RoPE bf16")
# No indexer on this call: top-k buffer must be untouched.
assert (topk == 7).all(), "topk buffer should be untouched (no indexer)"
# ── fused_q ──────────────────────────────────────────────────────────────────
@@ -400,6 +470,75 @@ def test_fused_q_no_indexer(num_tokens: int):
assert_fp8(mqa[:, :, KV_LORA:], (qpe_ref / s).to(FP8), "mqa q_pe")
@pytest.mark.parametrize("num_tokens", [1, 17, 512])
@pytest.mark.parametrize("has_indexer", [True, False])
def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
"""bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is
produced (bf16, unquantized); ql_nope is consumed directly by the caller."""
torch.manual_seed(6)
dev = "cuda"
max_pos = 8192
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
q_pe = torch.randn(
num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
)
ql_nope = torch.randn(
num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
)
q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
index_q = index_w = idx_cos_sin = None
if has_indexer:
index_q = torch.randn(
num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16
)
index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32)
idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
iq_fp8, iw_out, q_pe_out = K.fused_q(
pos,
q_pe,
q_cos_sin,
index_q,
idx_cos_sin,
ql_nope,
q_scale,
index_w,
INDEX_HEAD_DIM**-0.5,
INDEX_HEADS**-0.5,
has_indexer=has_indexer,
index_rope_interleave=False,
quantize_mqa=False,
)
# MQA query: only the RoPE'd q_pe, bf16, unquantized.
assert q_pe_out.dtype == torch.bfloat16
assert q_pe_out.shape == (num_tokens, NUM_HEADS, ROPE_DIM)
qpe_ref = rope(
q_pe.float(),
pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
q_cos_sin,
interleave=True,
)
assert_bf16(q_pe_out, qpe_ref, "bf16 q_pe RoPE")
# Indexer-Q is unchanged on this path (still UE8M0 fp8 + folded weights).
if has_indexer:
assert index_q is not None
iq_ref = rope(
index_q.float(),
pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS),
idx_cos_sin,
interleave=False,
)
q_ref, scale_ref = ue8m0_quant(iq_ref)
assert_fp8(iq_fp8, q_ref, "indexer-Q fp8 (bf16-query path)")
iw_ref = index_w * scale_ref * (INDEX_HEAD_DIM**-0.5) * (INDEX_HEADS**-0.5)
torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
# ── fused_eh_norm (MTP) ──────────────────────────────────────────────────────
+17 -2
View File
@@ -42,7 +42,11 @@ VLLM_RUNNER_BASE_KWARGS = {
}
def run_test(vllm_runner, audio_assets, lora_request, expected_suffix, **kwargs):
def run_test(
vllm_runner, audio_assets, monkeypatch, lora_request, expected_suffix, **kwargs
):
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
inputs = [([AUDIO_PROMPT], [audio_assets[0].audio_and_sample_rate[0]])]
# Apply any additional kwargs as overrides to the base kwargs
@@ -66,11 +70,13 @@ def run_test(vllm_runner, audio_assets, lora_request, expected_suffix, **kwargs)
def test_active_default_mm_lora(
vllm_runner: type[VllmRunner],
audio_assets: AudioTestAssets,
monkeypatch: pytest.MonkeyPatch,
):
"""Ensure that we can use the default audio lora."""
run_test(
vllm_runner,
audio_assets,
monkeypatch,
lora_request=None,
default_mm_loras={"audio": AUDIO_LORA_PATH},
expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
@@ -84,12 +90,14 @@ def test_active_default_mm_lora(
def test_inactive_default_mm_lora(
vllm_runner: type[VllmRunner],
audio_assets: AudioTestAssets,
monkeypatch: pytest.MonkeyPatch,
):
"""Ensure that modalities are filtered properly."""
# Default image lora won't be active since we only pass audio
run_test(
vllm_runner,
audio_assets,
monkeypatch,
lora_request=None,
default_mm_loras={"image": IMAGE_LORA_PATH},
expected_suffix=RESPONSE_SUFFIX_WITHOUT_LORA,
@@ -103,11 +111,13 @@ def test_inactive_default_mm_lora(
def test_default_mm_lora_succeeds_with_redundant_lora_request(
vllm_runner: type[VllmRunner],
audio_assets: AudioTestAssets,
monkeypatch: pytest.MonkeyPatch,
):
"""Ensure that redundantly providing the lora works."""
run_test(
vllm_runner,
audio_assets,
monkeypatch,
lora_request=LoRARequest("audio", 1, AUDIO_LORA_PATH),
default_mm_loras={"audio": AUDIO_LORA_PATH},
expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
@@ -121,12 +131,14 @@ def test_default_mm_lora_succeeds_with_redundant_lora_request(
def test_default_mm_lora_fails_with_overridden_lora_request(
vllm_runner: type[VllmRunner],
audio_assets: AudioTestAssets,
monkeypatch: pytest.MonkeyPatch,
):
"""Ensure that if the lora_request conflicts with default_mm_loras,
we use the lora_request."""
run_test(
vllm_runner,
audio_assets,
monkeypatch,
lora_request=LoRARequest("speech", 2, AUDIO_LORA_PATH),
default_mm_loras={"audio": IMAGE_LORA_PATH},
expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
@@ -134,7 +146,10 @@ def test_default_mm_lora_fails_with_overridden_lora_request(
@create_new_process_for_each_test()
def test_default_mm_lora_does_not_expand_string_reqs(vllm_runner):
def test_default_mm_lora_does_not_expand_string_reqs(vllm_runner, monkeypatch):
# See run_test: force spawn to avoid the forked-child CUDA re-init crash.
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
class MockEngineException(Exception):
pass
+21
View File
@@ -1,7 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import importlib.metadata
from importlib.util import find_spec
import pytest
import torch
from packaging import version
import vllm
from vllm.lora.request import LoRARequest
@@ -9,6 +14,22 @@ from vllm.platforms import current_platform
from ..utils import multi_gpu_test
# Require amd-quark >= 0.12 on torch >= 2.11.
# Earlier torch releases work with older quark versions. See
# https://github.com/amd/Quark/issues/34
# TODO: Remove once amd-quark>=0.12.0
QUARK_TORCH_COMPATIBLE = find_spec("quark") is not None and (
version.parse(importlib.metadata.version("amd-quark")) >= version.parse("0.12.0")
if version.parse(torch.__version__.split("+")[0]) >= version.parse("2.11")
else True
)
if current_platform.is_rocm() and not QUARK_TORCH_COMPATIBLE:
pytest.skip(
"This test requires amd-quark >= 0.12 on torch >= 2.11.",
allow_module_level=True,
)
MODEL_PATH = "openai/gpt-oss-20b"
PROMPT_TEMPLATE = """<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.
@@ -326,39 +326,6 @@ VLM_TEST_SETTINGS = {
large_gpu_mark(min_gb=64),
],
),
"aya_vision": VLMTestInfo(
models=["CohereLabs/aya-vision-8b"],
test_type=(VLMTestType.IMAGE),
prompt_formatter=lambda img_prompt: f"<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{img_prompt}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", # noqa: E501
single_image_prompts=IMAGE_ASSETS.prompts(
{
"stop_sign": "<image>What's the content in the center of the image?",
"cherry_blossom": "<image>What is the season?",
}
),
multi_image_prompt="<image><image>Describe the two images in detail.",
max_model_len=4096,
max_num_seqs=2,
auto_cls=AutoModelForImageTextToText,
vllm_runner_kwargs={"mm_processor_kwargs": {"crop_to_patches": True}},
),
"aya_vision-multi_image": VLMTestInfo(
models=["CohereLabs/aya-vision-8b"],
test_type=(VLMTestType.MULTI_IMAGE),
prompt_formatter=lambda img_prompt: f"<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{img_prompt}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", # noqa: E501
single_image_prompts=IMAGE_ASSETS.prompts(
{
"stop_sign": "<image>What's the content in the center of the image?",
"cherry_blossom": "<image>What is the season?",
}
),
multi_image_prompt="<image><image>Describe the two images in detail.",
max_model_len=4096,
max_num_seqs=2,
auto_cls=AutoModelForImageTextToText,
vllm_runner_kwargs={"mm_processor_kwargs": {"crop_to_patches": True}},
marks=[large_gpu_mark(min_gb=32)],
),
"blip2": VLMTestInfo(
models=["Salesforce/blip2-opt-2.7b"],
test_type=VLMTestType.IMAGE,
@@ -763,16 +730,6 @@ VLM_TEST_SETTINGS = {
auto_cls=AutoModelForImageTextToText,
vllm_output_post_proc=model_utils.llava_video_vllm_to_hf_output,
),
"mantis": VLMTestInfo(
models=["TIGER-Lab/Mantis-8B-siglip-llama3"],
test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
prompt_formatter=lambda img_prompt: f"<|start_header_id|>user<|end_header_id|>\n\n{img_prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n", # noqa: E501
max_model_len=4096,
get_stop_token_ids=lambda tok: [128009],
auto_cls=AutoModelForImageTextToText,
vllm_output_post_proc=model_utils.mantis_vllm_to_hf_output,
patch_hf_runner=model_utils.mantis_patch_hf_runner,
),
"minicpmv_25": VLMTestInfo(
models=["openbmb/MiniCPM-Llama3-V-2_5"],
test_type=VLMTestType.IMAGE,
@@ -1,152 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import os
import pytest
from tests.models.registry import HF_EXAMPLE_MODELS
from vllm import LLM, SamplingParams
MODEL_NAME = "nvidia/music-flamingo-2601-hf"
SINGLE_CONVERSATION = [
{
"role": "user",
"content": [
{
"type": "text",
"text": "Describe this track in full detail - tell me the "
"genre, tempo, and key, then dive into the instruments, "
"production style, and overall mood it creates.",
},
{
"type": "audio_url",
"audio_url": {
"url": "https://huggingface.co/datasets/nvidia/AudioSkills/"
"resolve/main/assets/song_1.mp3",
},
},
],
}
]
BATCHED_CONVERSATIONS = [
SINGLE_CONVERSATION,
[
{
"role": "user",
"content": [
{
"type": "text",
"text": "Generate a structured lyric sheet from the input music.",
},
{
"type": "audio_url",
"audio_url": {
"url": "https://huggingface.co/datasets/nvidia/"
"AudioSkills/resolve/main/assets/song_2.mp3",
},
},
],
}
],
]
def get_fixture_path(filename):
return os.path.join(
os.path.dirname(__file__), "../../fixtures/musicflamingo", filename
)
def load_expected_fixture(filename):
fixture_path = get_fixture_path(filename)
with open(fixture_path) as f:
return json.load(f)
def assert_output_matches(output, expected_text, expected_token_ids):
generated = output.outputs[0]
assert generated.text == expected_text
actual_token_ids = list(generated.token_ids)
assert (
actual_token_ids == expected_token_ids
or actual_token_ids == expected_token_ids[:-1]
or actual_token_ids[:-1] == expected_token_ids
)
@pytest.fixture(scope="module")
def llm():
model_info = HF_EXAMPLE_MODELS.get_hf_info("MusicFlamingoForConditionalGeneration")
model_info.check_transformers_version(on_fail="skip")
try:
llm = LLM(
model=MODEL_NAME,
dtype="bfloat16",
enforce_eager=True,
max_model_len=8192,
limit_mm_per_prompt={"audio": 1},
)
except Exception as e:
pytest.skip(f"Failed to load model {MODEL_NAME}: {e}")
# ROCm may compile decoder kernels on the first inference pass; warm up
# once so exact fixture assertions cover the steady-state path.
llm.chat(
messages=SINGLE_CONVERSATION,
sampling_params=SamplingParams(temperature=0.0, max_tokens=1),
use_tqdm=False,
)
return llm
def test_single_generation(llm):
expected = load_expected_fixture("expected_results_single.json")
outputs = llm.chat(
messages=SINGLE_CONVERSATION,
sampling_params=SamplingParams(temperature=0.0, max_tokens=50),
)
assert_output_matches(
outputs[0],
expected["transcriptions"][0],
expected["token_ids"][0],
)
def test_batched_generation(llm):
expected = load_expected_fixture("expected_results_batched.json")
outputs = llm.chat(
messages=BATCHED_CONVERSATIONS,
sampling_params=SamplingParams(temperature=0.0, max_tokens=50),
)
for i, output in enumerate(outputs):
assert_output_matches(
output,
expected["transcriptions"][i],
expected["token_ids"][i],
)
def test_single_and_batched_generation_match(llm):
sampling_params = SamplingParams(temperature=0.0, max_tokens=50)
single_output = llm.chat(
messages=SINGLE_CONVERSATION,
sampling_params=sampling_params,
)[0]
batched_output = llm.chat(
messages=BATCHED_CONVERSATIONS,
sampling_params=sampling_params,
)[0]
assert single_output.outputs[0].text == batched_output.outputs[0].text
assert list(single_output.outputs[0].token_ids) == list(
batched_output.outputs[0].token_ids
)
@@ -15,11 +15,14 @@ from mistral_common.tokens.tokenizers.tekken import SpecialTokenPolicy
from vllm import LLM, EngineArgs, SamplingParams
from vllm.assets.audio import AudioAsset
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.utils.math_utils import cdiv
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.v1.kv_cache_interface import SlidingWindowSpec
from ....utils import ROCM_ENGINE_KWARGS
MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
AUDIO_LAYER_NAME = "whisper_encoder.whisper_encoder.layers.0.layers.self_attn.attn"
ENGINE_CONFIG = {
"model": MODEL_NAME,
"max_model_len": 8192,
@@ -60,6 +63,31 @@ def _normalize(texts: list[str]) -> list[str]:
return texts
def assert_encoder_kv_cache_spec(engine: LLM) -> None:
vllm_config = engine.llm_engine.vllm_config
audio_config = vllm_config.model_config.hf_config.audio_config
kv_cache_specs_per_rank = engine.llm_engine.model_executor.get_kv_cache_specs()
assert len(kv_cache_specs_per_rank) == 1
kv_cache_specs = kv_cache_specs_per_rank[0]
assert AUDIO_LAYER_NAME in kv_cache_specs, kv_cache_specs.keys()
spec = kv_cache_specs[AUDIO_LAYER_NAME]
assert audio_config.sliding_window == 750
assert audio_config.block_pool_size == 4
assert isinstance(spec, SlidingWindowSpec)
assert spec.block_size == 16
assert spec.num_kv_heads == 128
assert spec.sliding_window == cdiv(750, 4) == 188
assert (
spec.max_admission_blocks_per_request(
max_num_batched_tokens=1,
max_model_len=vllm_config.model_config.max_model_len,
)
== 13
)
@pytest.fixture
def audio_assets() -> list[AudioAsset]:
return [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
@@ -71,7 +99,9 @@ def tokenizer() -> MistralTokenizer:
@pytest.fixture
def engine():
def engine(monkeypatch: pytest.MonkeyPatch):
# Disable multiprocessing allows us to access model executor from LLM engine
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
engine_args = EngineArgs(**ENGINE_CONFIG)
llm = LLM.from_engine_args(engine_args)
try:
@@ -95,6 +125,7 @@ async def async_engine():
def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
assert_encoder_kv_cache_spec(engine)
audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
def from_file(file_path: str):
@@ -167,15 +167,6 @@ def llava_onevision_vllm_to_hf_output(
return hf_output_ids, hf_output_str, out_logprobs
def mantis_vllm_to_hf_output(vllm_output: RunnerOutput, model: str) -> RunnerOutput:
"""Sanitize vllm output [mantis] to compare with hf output."""
output_ids, output_str, out_logprobs = vllm_output
hf_output_str = output_str + "<|eot_id|>"
return output_ids, hf_output_str, out_logprobs
def phi3v_vllm_to_hf_output(vllm_output: RunnerOutput, model: str) -> RunnerOutput:
"""Sanitize vllm output [phi3v] to be comparable with hf output."""
_, output_str, out_logprobs = vllm_output
@@ -940,29 +931,6 @@ def _internvl_generate(
return outputs
def mantis_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
from mantis.models.mllava import MLlavaProcessor
hf_model.processor = MLlavaProcessor.from_pretrained(hf_model.model_name)
orig_generate = hf_model.model.generate
tokenizer = hf_model.processor.tokenizer
def _generate(self, *args, **kwargs):
return orig_generate(
*args,
**kwargs,
eos_token_id=[
tokenizer.eos_token_id,
tokenizer.convert_tokens_to_ids("<|eot_id|>"),
],
)
hf_model.model.generate = types.MethodType(_generate, hf_model.model)
return hf_model
def minicpmv_25_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
orig_generate = hf_model.model.generate
@@ -1,228 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Copyright 2026 The vLLM team.
# Copyright 2026 NVIDIA CORPORATION and the HuggingFace Inc. team. All rights
# reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from importlib.metadata import version
from unittest.mock import MagicMock
import numpy as np
import pytest
import torch
from packaging.version import Version
from transformers import PretrainedConfig
from tests.models.registry import HF_EXAMPLE_MODELS
class MockMusicFlamingoConfig(PretrainedConfig):
model_type = "musicflamingo"
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.audio_config = PretrainedConfig()
self.text_config = PretrainedConfig()
class MockMusicFlamingoProcessor:
def __init__(self):
self.audio_token = "<sound>"
self.audio_token_id = 12345
self.audio_bos_token = "<|sound_bos|>"
self.audio_bos_token_id = 12346
self.audio_eos_token = "<|sound_eos|>"
self.audio_eos_token_id = 12347
self.max_audio_len = 1200
self.feature_extractor = MockFeatureExtractor()
def __call__(self, text=None, audio=None, **kwargs):
return {
"input_ids": torch.tensor([[1, 2, 3]], dtype=torch.long),
"input_features": torch.zeros((3, 80, 3000)),
"input_features_mask": torch.ones((3, 3000), dtype=torch.long),
}
class MockFeatureExtractor:
def __init__(self):
self.sampling_rate = 16000
self.chunk_length = 30
@pytest.fixture
def mock_ctx():
config = MockMusicFlamingoConfig()
ctx = MagicMock()
ctx.get_hf_config.return_value = config
ctx.get_hf_processor.return_value = MockMusicFlamingoProcessor()
ctx.call_hf_processor.side_effect = lambda processor, data, kwargs: processor(
**data, **kwargs
)
ctx.model_config.hf_config = config
return ctx
@pytest.fixture(autouse=True)
def check_transformers_version():
model_info = HF_EXAMPLE_MODELS.get_hf_info("MusicFlamingoForConditionalGeneration")
model_info.check_transformers_version(on_fail="skip")
def test_musicflamingo_chunk_counting_without_rote_timestamps(mock_ctx):
from vllm.model_executor.models.musicflamingo import (
MusicFlamingoDummyInputsBuilder,
MusicFlamingoMultiModalProcessor,
MusicFlamingoProcessingInfo,
)
info = MusicFlamingoProcessingInfo(mock_ctx)
processor = MusicFlamingoMultiModalProcessor(
info, MusicFlamingoDummyInputsBuilder(info)
)
sr = 16000
audio_1 = np.zeros(30 * sr)
audio_2 = np.zeros(45 * sr)
mm_data = {"audio": [audio_1, audio_2]}
prompt = "<|user|>Listen.<|end|>"
processed = processor._call_hf_processor(prompt, mm_data, {}, {})
chunk_counts = processed["chunk_counts"]
assert chunk_counts.tolist() == [1, 2]
assert "rote_timestamps" not in processed
assert processed["feature_attention_mask"].shape == (3, 3000)
def test_musicflamingo_dummy_text_uses_plain_audio_tokens(mock_ctx):
from vllm.model_executor.models.musicflamingo import (
MusicFlamingoDummyInputsBuilder,
MusicFlamingoProcessingInfo,
)
info = MusicFlamingoProcessingInfo(mock_ctx)
builder = MusicFlamingoDummyInputsBuilder(info)
assert builder.get_dummy_text({"audio": 2}) == "<sound><sound>"
@pytest.mark.skipif(
Version(version("transformers")) >= Version("5.5"),
reason="transformers v5.5 added native MusicFlamingoForConditionalGeneration "
"with a different get_audio_features signature (requires input_ids)",
)
def test_musicflamingo_audio_feature_pipeline_matches_hf_small_config():
from transformers.models.musicflamingo import (
modeling_musicflamingo as hf_musicflamingo_modeling,
)
from transformers.models.musicflamingo.configuration_musicflamingo import (
MusicFlamingoConfig,
)
from vllm.model_executor.models.audioflamingo3 import (
_build_audio_encoder_attention_mask,
_flatten_valid_audio_embeddings,
)
from vllm.model_executor.models.musicflamingo import (
MusicFlamingoEncoder,
MusicFlamingoMultiModalProjector,
MusicFlamingoRotaryEmbedding,
apply_rotary_time_emb,
)
text_config = {
"model_type": "qwen2",
"intermediate_size": 64,
"initializer_range": 0.02,
"hidden_size": 32,
"max_position_embeddings": 1024,
"num_hidden_layers": 2,
"num_attention_heads": 4,
"num_key_value_heads": 2,
"vocab_size": 128,
"pad_token_id": 1,
"use_mrope": False,
}
audio_config = {
"hidden_size": 16,
"num_attention_heads": 4,
"intermediate_size": 32,
"num_hidden_layers": 2,
"num_mel_bins": 80,
"max_source_positions": 1500,
"dropout": 0.0,
"attention_dropout": 0.0,
"activation_dropout": 0.0,
"encoder_layerdrop": 0.0,
}
torch.manual_seed(0)
config = MusicFlamingoConfig(
text_config=text_config,
audio_config=audio_config,
audio_token_id=0,
head_dim=8,
rope_parameters={"rope_type": "default", "rope_theta": 2048},
)
hf_model = hf_musicflamingo_modeling.MusicFlamingoForConditionalGeneration(
config
).eval()
vllm_encoder = MusicFlamingoEncoder(config.audio_config).eval()
vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict())
vllm_projector = MusicFlamingoMultiModalProjector(config).eval()
vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict())
vllm_rope = MusicFlamingoRotaryEmbedding(config).eval()
vllm_rope.load_state_dict(hf_model.pos_emb.state_dict(), strict=False)
input_features = torch.randn(3, 80, 3000)
feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool)
feature_attention_mask[0, :3000] = True
feature_attention_mask[1, :2500] = True
feature_attention_mask[2, :1500] = True
rote_timestamps = (
torch.arange(750, dtype=torch.float32).unsqueeze(0).repeat(3, 1) * 0.04
)
hf_output = hf_model.get_audio_features(
input_features,
feature_attention_mask,
rote_timestamps=rote_timestamps,
return_dict=True,
).pooler_output
vllm_attention_mask = _build_audio_encoder_attention_mask(
feature_attention_mask,
dtype=vllm_encoder.conv1.weight.dtype,
device=vllm_encoder.conv1.weight.device,
)
vllm_hidden_states = vllm_encoder(
input_features,
attention_mask=vllm_attention_mask,
)
cos, sin = vllm_rope(rote_timestamps, seq_len=vllm_hidden_states.shape[-2])
vllm_hidden_states = apply_rotary_time_emb(vllm_hidden_states, cos, sin)
vllm_output, _ = _flatten_valid_audio_embeddings(
vllm_projector(vllm_hidden_states),
feature_attention_mask,
)
torch.testing.assert_close(vllm_output, hf_output)
+1 -15
View File
@@ -780,14 +780,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"vllm": "Needs https://github.com/huggingface/transformers/pull/43538"
},
),
"MusicFlamingoForConditionalGeneration": _HfExamplesInfo(
"nvidia/music-flamingo-2601-hf",
min_transformers_version="5.5.0",
transformers_version_reason={
"vllm": "Needs https://github.com/huggingface/transformers/pull/43538"
},
),
"AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereLabs/aya-vision-8b"),
"BagelForConditionalGeneration": _HfExamplesInfo("ByteDance-Seed/BAGEL-7B-MoT"),
"BeeForConditionalGeneration": _HfExamplesInfo(
"Open-Bee/Bee-8B-RL",
@@ -1059,12 +1051,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"LlavaOnevisionForConditionalGeneration": _HfExamplesInfo(
"llava-hf/llava-onevision-qwen2-0.5b-ov-hf"
),
"MantisForConditionalGeneration": _HfExamplesInfo(
"TIGER-Lab/Mantis-8B-siglip-llama3",
max_transformers_version="4.48",
transformers_version_reason={"hf": "HF model is not compatible."},
hf_overrides={"architectures": ["MantisForConditionalGeneration"]},
),
"MiDashengLMModel": _HfExamplesInfo(
"mispeech/midashenglm-7b", trust_remote_code=True
),
@@ -1408,7 +1394,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
# [DFlash]
"DFlashDraftModel": _HfExamplesInfo(
"Qwen/Qwen3.5-4B",
speculative_model="z-lab/Qwen3.5-4B-DFlash",
speculative_model="z-lab/Qwen3-4B-DFlash-b16",
use_original_num_layers=True, # Need all layers since DFlash has >1 layer,
max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env
max_num_seqs=32,
+19 -1
View File
@@ -10,7 +10,11 @@ import pytest
from PIL import Image
from vllm.assets.base import get_vllm_public_assets
from vllm.assets.video import video_to_ndarrays, video_to_pil_images_list
from vllm.assets.video import (
video_get_metadata,
video_to_ndarrays,
video_to_pil_images_list,
)
from vllm.multimodal.media import ImageMediaIO, VideoMediaIO
from vllm.multimodal.video import VIDEO_LOADER_REGISTRY, VideoLoader
@@ -112,6 +116,20 @@ def test_opencv_video_io_colorspace(tmp_path, is_color: bool, fourcc: str, ext:
assert np.nanmean(sim) > 0.99
def test_opencv_video_metadata_matches_sampled_frame_timeline(tmp_path):
image_path = f"{tmp_path}/test_metadata_image.png"
Image.new("RGB", (8, 8), color=(255, 0, 0)).save(image_path)
video_path = f"{tmp_path}/test_metadata_video.mp4"
create_video_from_image(image_path, video_path, num_frames=10, fps=5.0)
metadata = video_get_metadata(video_path, num_frames=4)
assert metadata["fps"] == pytest.approx(5.0)
assert metadata["duration"] == pytest.approx(2.0)
assert metadata["frames_indices"] == [0, 3, 6, 9]
assert metadata["total_num_frames"] == 4
NUM_FRAMES = 10
FAKE_OUTPUT_1 = np.random.rand(NUM_FRAMES, 1280, 720, 3)
FAKE_OUTPUT_2 = np.random.rand(NUM_FRAMES, 1280, 720, 3)

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