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
d8f483dc30 [Spec Decode] Fix hidden-state extraction block size for hybrid verifiers (#46301)
Signed-off-by: Igor Margulis <igor.margulis@intel.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: mgoin <mgoin64@gmail.com>
2026-06-30 08:19:51 -07:00
Nicolò LucchesiandGitHub dc148dc4d7 [CI][Bugfix] Fix Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) (#47157)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
2026-06-30 23:14:13 +08:00
tc-mbandGitHub 7cf7cbcd95 [Bugfix] MiniCPM-V 4.6: fix grid rows/cols swap in placeholder generation (#45918)
Signed-off-by: tc-mb <tianchi_cai@icloud.com>
2026-06-30 08:12:44 -07:00
Juan Pérez de AlgabaGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
c231d1f290 fix(security): bound tokenizer work when explicit truncation_side is set (#47007)
Signed-off-by: jperezde <jperezde@redhat.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-30 23:08:51 +08:00
Giancarlo DelfinandGitHub db808b3961 [Model Runner V2][Spec Decode] Implement block verification for rejection sampling (#46781)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
2026-06-30 08:07:24 -07:00
Arsalan ShakilandGitHub 00ebf19cca [Bugfix][Quant] Raise actionable error instead of bare assert for group-size/TP mismatch (#46230) (#46236)
Signed-off-by: Arsalan Shakil <shakil.arsalan@yahoo.com>
2026-06-30 14:57:14 +00:00
Seiji EicherGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
ded6676458 [Bugfix] Seed RayExecutorV2 TCPStore port by DP rank to avoid collisions (#45960)
Signed-off-by: Seiji Eicher <seiji@anyscale.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-30 07:37:34 -07:00
Bugen ZhaoandGitHub 7a327f0b4f [Rust Frontend] Simplify unit tests with shared TestTokenizer (#47125)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-30 15:34:43 +01:00
Harry MellorandGitHub 1ab9522935 Remove more unnecessary load_weights methods (#47058)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-06-30 15:22:16 +01:00
0fc2512094 [KV Offload] Pass ScheduleEndContext to on_schedule_end hook (#46450)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
2026-06-30 17:07:12 +03:00
Harry MellorandGitHub 62c7d8009f Forward fix nightly errors from #44589 (#47151)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-06-30 14:02:34 +00:00
Isotr0pyandGitHub ab80b3dff4 [CI/Build] Bump PyNvVideoCodec version (#47139) 2026-06-30 06:38:46 -07:00
Qiming ZhangandGitHub 91055efd36 [XPU] C++ implementation for get_memory_info (#47134)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
2026-06-30 21:34:47 +08:00
Bugen ZhaoandGitHub 3675bcff67 [Rust Frontend] Refactor TLS serve path with unified MaybeTlsListener (#47101)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-30 14:31:58 +01:00
Bugen ZhaoandGitHub bdbd7278b6 [Rust Frontend] Extend renderer/parser roundtrip tests to support token ids (#47110)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-30 14:27:45 +01:00
Harry MellorandGitHub 5dc36a4fa5 [Model] Remove Tarsier, Tarsier2 (#47143)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-06-30 13:20:33 +00:00
aab7af0bcb [Bugfix][ROCm][MLA] Pass q/kv dtypes to get_mla_metadata_v1 in FP8 decode (#46997)
Signed-off-by: pei.zhang <pei.zhang@amd.com>
Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
2026-06-30 05:31:16 -07:00
377 changed files with 11584 additions and 11634 deletions
+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}" \
+35 -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
@@ -2828,9 +2816,10 @@ steps:
- rocm-smi
- python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
- pytest -v -s tests/kernels/attention/test_attention_selector.py
- 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
@@ -2979,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
@@ -2993,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
@@ -3019,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
@@ -3034,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
@@ -3177,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
@@ -3189,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);
+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 -10
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` | ✅︎ | ✅︎ |
@@ -626,8 +624,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `Step3VLForConditionalGeneration` | Step3-VL | T + I<sup>+</sup> | `stepfun-ai/step3` | | ✅︎ |
| `StepVLForConditionalGeneration` | Step3-VL-10B | T + I<sup>+</sup> | `stepfun-ai/Step3-VL-10B` | | ✅︎ |
| `Step3p7ForConditionalGeneration` | Step-3.7-Flash | T + I<sup>+</sup> | `stepfun-ai/Step-3.7-Flash` | | ✅︎ |
| `TarsierForConditionalGeneration` | Tarsier | T + I<sup>E+</sup> | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ |
| `Tarsier2ForConditionalGeneration`<sup>^</sup> | Tarsier2 | T + I<sup>E+</sup> + V<sup>E+</sup> | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ |
| `UltravoxModel` | Ultravox | T + A<sup>E+</sup> | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ |
| `UnlimitedOCRForCausalLM` | Unlimited-OCR | T + I<sup>+</sup> | `baidu/Unlimited-OCR`, etc. | ✅︎ | ✅︎ |
@@ -680,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"
@@ -1275,55 +1242,6 @@ def load_step_vl(question: str, image_urls: list[str]) -> ModelRequestData:
)
def load_tarsier(question: str, image_urls: list[str]) -> ModelRequestData:
model_name = "omni-research/Tarsier-7b"
engine_args = EngineArgs(
model=model_name,
trust_remote_code=True,
max_model_len=4096,
limit_mm_per_prompt={"image": len(image_urls)},
)
prompt = f"USER: {'<image>' * len(image_urls)}\n{question}\n ASSISTANT:"
image_data = [fetch_image(url) for url in image_urls]
return ModelRequestData(
engine_args=engine_args,
prompt=prompt,
image_data=image_data,
)
def load_tarsier2(question: str, image_urls: list[str]) -> ModelRequestData:
model_name = "omni-research/Tarsier2-Recap-7b"
engine_args = EngineArgs(
model=model_name,
trust_remote_code=True,
max_model_len=32768,
limit_mm_per_prompt={"image": len(image_urls)},
hf_overrides={
"architectures": ["Tarsier2ForConditionalGeneration"],
"model_type": "tarsier2",
},
)
prompt = (
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
f"<|im_start|>user\n<|vision_start|>{'<|image_pad|>' * len(image_urls)}"
f"<|vision_end|>{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
image_data = [fetch_image(url) for url in image_urls]
return ModelRequestData(
engine_args=engine_args,
prompt=prompt,
image_data=image_data,
)
# GLM-4.1V
def load_glm4_1v(question: str, image_urls: list[str]) -> ModelRequestData:
model_name = "zai-org/GLM-4.1V-9B-Thinking"
@@ -1469,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,
@@ -1507,8 +1424,6 @@ model_example_map = {
"smolvlm": load_smolvlm,
"step3": load_step3,
"stepvl": load_step_vl,
"tarsier": load_tarsier,
"tarsier2": load_tarsier2,
"glm4_1v": load_glm4_1v,
"glm4_5v": load_glm4_5v,
"glm4_5v_fp8": load_glm4_5v_fp8,
@@ -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"]
@@ -2347,68 +2303,8 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData:
)
# omni-research/Tarsier-7b
def run_tarsier(questions: list[str], modality: str) -> ModelRequestData:
assert modality == "image"
model_name = "omni-research/Tarsier-7b"
engine_args = EngineArgs(
model=model_name,
trust_remote_code=True,
max_model_len=4096,
limit_mm_per_prompt={modality: 1},
)
prompts = [(f"USER: <image>\n{question} ASSISTANT:") for question in questions]
return ModelRequestData(
engine_args=engine_args,
prompts=prompts,
)
def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData:
model_name = "omni-research/Tarsier2-Recap-7b"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
hf_overrides={
"architectures": ["Tarsier2ForConditionalGeneration"],
"model_type": "tarsier2",
},
limit_mm_per_prompt=mm_limit,
)
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
if modality == "image":
placeholder = image_placeholder
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
prompts = [
(
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
f"<|im_start|>user\n{placeholder}"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
for question in questions
]
return ModelRequestData(
engine_args=engine_args,
prompts=prompts,
)
model_example_map = {
"aria": run_aria,
"aya_vision": run_aya_vision,
"bagel": run_bagel,
"cheers": run_cheers,
"bee": run_bee,
@@ -2449,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,
@@ -2479,8 +2374,6 @@ model_example_map = {
"smolvlm": run_smolvlm,
"step3": run_step3,
"stepvl": run_step_vl,
"tarsier": run_tarsier,
"tarsier2": run_tarsier2,
}
+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))
+1 -1
View File
@@ -8,7 +8,7 @@ torch==2.11.0
torchaudio==2.11.0
# These must be updated alongside torch
torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
PyNvVideoCodec==2.0.4
PyNvVideoCodec==2.1.0
# FlashInfer should be updated together with the Dockerfile
flashinfer-python==0.6.13
flashinfer-cubin==0.6.13
+17 -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"
@@ -5310,6 +5325,7 @@ dependencies = [
"vllm-llm",
"vllm-metrics",
"vllm-text",
"vllm-tokenizer",
"zeromq",
]
@@ -5323,7 +5339,6 @@ dependencies = [
"enum-as-inner",
"expect-test",
"futures",
"hf-hub",
"itertools 0.14.0",
"reqwest",
"serde",
@@ -5338,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" }
+2
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
@@ -50,6 +51,7 @@ tokio.workspace = true
tracing-subscriber.workspace = true
uuid.workspace = true
vllm-engine-core-client = { workspace = true, features = ["test-util"] }
vllm-tokenizer = { workspace = true, features = ["test-utils"] }
zeromq.workspace = true
[lints]
+3 -26
View File
@@ -154,7 +154,8 @@ mod tests {
use thiserror_ext::AsReport as _;
use vllm_text::Prompt;
use vllm_text::backend::hf::TokenizerSource;
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
use vllm_text::tokenizer::DynTokenizer;
use vllm_tokenizer::test_utils::TestTokenizer;
use super::HfChatBackend;
use crate::backend::{ChatBackend, LoadModelBackendsOptions, NewChatOutputProcessorOptions};
@@ -196,32 +197,8 @@ mod tests {
}
}
struct TestTokenizer;
impl Tokenizer for TestTokenizer {
fn encode(
&self,
_text: &str,
_add_special_tokens: bool,
) -> vllm_text::tokenizer::Result<Vec<u32>> {
Ok(Vec::new())
}
fn decode(
&self,
_token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_text::tokenizer::Result<String> {
Ok(String::new())
}
fn token_to_id(&self, _token: &str) -> Option<u32> {
None
}
}
fn test_tokenizer() -> DynTokenizer {
Arc::new(TestTokenizer)
Arc::new(TestTokenizer::new())
}
fn backend_for_selection(
+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),
}
+10 -56
View File
@@ -563,7 +563,7 @@ mod tests {
use llm_multimodal::TokenId;
use vllm_engine_core_client::protocol::tensor::WireArrayData;
use vllm_text::tokenizer::{IncrementalDecoder, Tokenizer, TokenizerError};
use vllm_tokenizer::test_utils::TestTokenizer;
use super::*;
@@ -574,60 +574,14 @@ mod tests {
const LLAMA4_TILE_X_SEPARATOR_ID: u32 = 200093;
const LLAMA4_TILE_Y_SEPARATOR_ID: u32 = 200094;
struct TestTokenizer;
impl Tokenizer for TestTokenizer {
fn encode(
&self,
text: &str,
_add_special_tokens: bool,
) -> std::result::Result<Vec<u32>, TokenizerError> {
Ok(match text {
"<|image|>" => vec![LLAMA4_IMAGE_ID],
text => text.bytes().map(u32::from).collect(),
})
}
fn decode(
&self,
_token_ids: &[u32],
_skip_special_tokens: bool,
) -> std::result::Result<String, TokenizerError> {
Ok(String::new())
}
fn token_to_id(&self, token: &str) -> Option<u32> {
match token {
"<|image_start|>" => Some(LLAMA4_IMAGE_START_ID),
"<|image_end|>" => Some(LLAMA4_IMAGE_END_ID),
"<|image|>" => Some(LLAMA4_IMAGE_ID),
"<|patch|>" => Some(LLAMA4_PATCH_ID),
"<|tile_x_separator|>" => Some(LLAMA4_TILE_X_SEPARATOR_ID),
"<|tile_y_separator|>" => Some(LLAMA4_TILE_Y_SEPARATOR_ID),
_ => None,
}
}
fn id_to_token(&self, id: u32) -> Option<String> {
match id {
LLAMA4_IMAGE_START_ID => Some("<|image_start|>".to_string()),
LLAMA4_IMAGE_END_ID => Some("<|image_end|>".to_string()),
LLAMA4_IMAGE_ID => Some("<|image|>".to_string()),
LLAMA4_PATCH_ID => Some("<|patch|>".to_string()),
LLAMA4_TILE_X_SEPARATOR_ID => Some("<|tile_x_separator|>".to_string()),
LLAMA4_TILE_Y_SEPARATOR_ID => Some("<|tile_y_separator|>".to_string()),
_ => None,
}
}
fn create_decode_stream(
&self,
_prompt_token_ids: &[u32],
_skip_special_tokens: bool,
_min_bytes_to_buffer: usize,
) -> Box<dyn IncrementalDecoder + '_> {
unreachable!("not used")
}
fn llama4_tokenizer() -> TestTokenizer {
TestTokenizer::new()
.with_regular_token("<|image_start|>", LLAMA4_IMAGE_START_ID)
.with_regular_token("<|image_end|>", LLAMA4_IMAGE_END_ID)
.with_regular_token("<|image|>", LLAMA4_IMAGE_ID)
.with_regular_token("<|patch|>", LLAMA4_PATCH_ID)
.with_regular_token("<|tile_x_separator|>", LLAMA4_TILE_X_SEPARATOR_ID)
.with_regular_token("<|tile_y_separator|>", LLAMA4_TILE_Y_SEPARATOR_ID)
}
fn test_info(model_type: &str, config: serde_json::Value) -> MultimodalModelInfo {
@@ -635,7 +589,7 @@ mod tests {
model_id: format!("{model_type}-test"),
model_type: Some(model_type.to_string()),
config,
tokenizer: TokenizerResolver(Arc::new(TestTokenizer)),
tokenizer: TokenizerResolver(Arc::new(llama4_tokenizer())),
};
let spec = context
.resolve_model_spec()
+7 -34
View File
@@ -189,46 +189,19 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor {
mod tests {
use std::sync::Arc;
use vllm_tokenizer::Tokenizer;
use vllm_tokenizer::test_utils::TestTokenizer;
use super::DefaultChatOutputProcessor;
use crate::Error;
use crate::parser::ParserSelection;
use crate::request::ChatRequest;
struct FakeTokenizer;
impl Tokenizer for FakeTokenizer {
fn encode(
&self,
text: &str,
_add_special_tokens: bool,
) -> vllm_tokenizer::Result<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok(token_ids
.iter()
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
.collect())
}
fn token_to_id(&self, token: &str) -> Option<u32> {
match token {
"<|channel>" => Some(1),
"<channel|>" => Some(2),
_ => None,
}
}
}
fn tokenizer() -> Arc<FakeTokenizer> {
Arc::new(FakeTokenizer)
fn tokenizer() -> Arc<TestTokenizer> {
Arc::new(
TestTokenizer::new()
.with_regular_token("<|channel>", 256)
.with_regular_token("<channel|>", 257),
)
}
#[test]
-14
View File
@@ -1,14 +1,8 @@
//! Harmony output tests share the upstream `openai-harmony` tiktoken cache.
//!
//! Use a file lock for tests that load the encoding so `cargo nextest` cannot
//! start multiple processes that concurrently populate the same cache file.
use std::sync::Arc;
use futures::executor::block_on;
use futures::{TryStreamExt as _, stream};
use openai_harmony::chat::{Message, Role};
use serial_test::file_serial;
use vllm_text::output::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, Finished};
use super::*;
@@ -91,7 +85,6 @@ fn request_with_tools() -> ChatRequest {
}
#[test]
#[file_serial(harmony_tiktoken_cache)]
fn interrupted_final_message_is_preserved() {
let tokens = completion_tokens(&[text_message("final", "hello")]);
let events = block_on(collect_events(
@@ -127,7 +120,6 @@ fn interrupted_final_message_is_preserved() {
}
#[test]
#[file_serial(harmony_tiktoken_cache)]
fn eos_flush_preserves_trailing_replacement_text() {
let mut tokens = completion_tokens(&[text_message("final", "Hi")]);
tokens.pop();
@@ -153,7 +145,6 @@ fn eos_flush_preserves_trailing_replacement_text() {
}
#[test]
#[file_serial(harmony_tiktoken_cache)]
fn interrupted_analysis_message_is_preserved() {
let tokens = completion_tokens(&[text_message("analysis", "think")]);
let events = block_on(collect_events(
@@ -189,7 +180,6 @@ fn interrupted_analysis_message_is_preserved() {
}
#[test]
#[file_serial(harmony_tiktoken_cache)]
fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() {
let tokens = completion_tokens(&[
text_message("commentary", "Let me check."),
@@ -217,7 +207,6 @@ fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() {
}
#[test]
#[file_serial(harmony_tiktoken_cache)]
fn multiple_messages_get_newline_separators() {
let tokens = completion_tokens(&[
text_message("analysis", "first think"),
@@ -249,7 +238,6 @@ fn multiple_messages_get_newline_separators() {
}
#[test]
#[file_serial(harmony_tiktoken_cache)]
fn tool_calls_stream_arguments_and_finish_with_local_id_shape() {
let tokens = completion_tokens(&[tool_message(
"get_weather",
@@ -302,7 +290,6 @@ fn tool_calls_stream_arguments_and_finish_with_local_id_shape() {
}
#[test]
#[file_serial(harmony_tiktoken_cache)]
fn semantic_events_precede_same_update_logprobs() {
let tokens = completion_tokens(&[text_message("final", "hello")]);
let events = block_on(collect_events(
@@ -353,7 +340,6 @@ fn rejects_generic_parser_overrides() {
}
#[test]
#[file_serial(harmony_tiktoken_cache)]
fn allows_auto_auto_only() {
validate_harmony_parser_overrides(&ParserSelection::Auto, &ParserSelection::Auto).unwrap();
let _ = HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap();
+2 -25
View File
@@ -1,32 +1,9 @@
use std::sync::Arc;
use vllm_tokenizer::Tokenizer;
use vllm_tokenizer::test_utils::TestTokenizer;
use super::{ReasoningParserFactory, names};
struct FakeTokenizer;
impl Tokenizer for FakeTokenizer {
fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok(token_ids
.iter()
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
.collect())
}
fn token_to_id(&self, _token: &str) -> Option<u32> {
None
}
}
#[test]
fn factory_contains_and_lists_registered_parsers() {
let factory = ReasoningParserFactory::new();
@@ -107,7 +84,7 @@ fn factory_resolves_minimax_m3_before_generic_minimax() {
#[test]
fn factory_rejects_unknown_parser_names() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(TestTokenizer::new());
let factory = ReasoningParserFactory::new();
let error = match factory.create("missing", tokenizer) {
Ok(_) => panic!("expected parser lookup to fail"),
+6 -31
View File
@@ -75,39 +75,14 @@ impl UnifiedParserFactory {
mod tests {
use std::sync::Arc;
use vllm_tokenizer::Tokenizer;
use vllm_tokenizer::test_utils::TestTokenizer;
use super::{UnifiedParserFactory, names};
struct FakeTokenizer;
impl Tokenizer for FakeTokenizer {
fn encode(
&self,
text: &str,
_add_special_tokens: bool,
) -> vllm_tokenizer::Result<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok(token_ids
.iter()
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
.collect())
}
fn token_to_id(&self, token: &str) -> Option<u32> {
match token {
"<|channel>" => Some(1),
"<channel|>" => Some(2),
_ => None,
}
}
fn tokenizer() -> TestTokenizer {
TestTokenizer::new()
.with_regular_token("<|channel>", 256)
.with_regular_token("<channel|>", 257)
}
#[test]
@@ -119,6 +94,6 @@ mod tests {
factory.resolve_name_for_model("google/gemma-4-27b-it"),
Some(names::GEMMA4)
);
factory.create(names::GEMMA4, &[], Arc::new(FakeTokenizer)).unwrap();
factory.create(names::GEMMA4, &[], Arc::new(tokenizer())).unwrap();
}
}
+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()
+34 -103
View File
@@ -21,15 +21,17 @@ use vllm_engine_core_client::protocol::{
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
use vllm_llm::Llm;
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
use vllm_text::tokenizer::DynTokenizer;
use vllm_text::{
DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, DecodedTokenLogprob, Prompt,
TextBackend,
};
use vllm_tokenizer::test_utils::TestTokenizer;
use zeromq::prelude::{SocketRecv, SocketSend};
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
const SPECIAL_STOP_TOKEN_ID: u32 = 256;
const UNKNOWN_DECODE_TOKEN_ID: u32 = 10_000;
fn request_output(
request_id: &str,
@@ -158,45 +160,18 @@ async fn connect_chat_llm_with_ipc(
struct FakeChatBackend {
has_template: bool,
model_id: String,
tokenizer: DynTokenizer,
}
#[derive(Debug)]
struct FakeChatTokenizer;
impl Tokenizer for FakeChatTokenizer {
fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
Ok(text.bytes().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
let bytes = token_ids
.iter()
.filter_map(|id| {
if skip_special_tokens && *id == SPECIAL_STOP_TOKEN_ID {
None
} else {
Some(*id as u8)
}
})
.collect::<Vec<_>>();
Ok(String::from_utf8_lossy(&bytes).into_owned())
}
fn token_to_id(&self, token: &str) -> Option<u32> {
match token {
"<think>" => Some(0xF001),
"</think>" => Some(0xF002),
"<|START_THINKING|>" => Some(0xF003),
"<|END_THINKING|>" => Some(0xF004),
"◁think▷" => Some(0xF005),
"◁/think▷" => Some(0xF006),
_ => None,
}
}
fn fake_chat_tokenizer() -> TestTokenizer {
TestTokenizer::new()
.with_special_token("<stop>", SPECIAL_STOP_TOKEN_ID)
.with_regular_token("<think>", 0xF001)
.with_regular_token("</think>", 0xF002)
.with_regular_token("<|START_THINKING|>", 0xF003)
.with_regular_token("<|END_THINKING|>", 0xF004)
.with_regular_token("◁think▷", 0xF005)
.with_regular_token("◁/think▷", 0xF006)
}
impl fmt::Debug for FakeChatBackend {
@@ -210,6 +185,7 @@ impl FakeChatBackend {
Self {
has_template: true,
model_id: "test-model".to_string(),
tokenizer: Arc::new(fake_chat_tokenizer()),
}
}
@@ -217,6 +193,7 @@ impl FakeChatBackend {
Self {
has_template: false,
model_id: "test-model".to_string(),
tokenizer: Arc::new(fake_chat_tokenizer()),
}
}
@@ -224,13 +201,19 @@ impl FakeChatBackend {
Self {
has_template: true,
model_id: model_id.into(),
tokenizer: Arc::new(fake_chat_tokenizer()),
}
}
fn with_tokenizer(mut self, tokenizer: DynTokenizer) -> Self {
self.tokenizer = tokenizer;
self
}
}
impl TextBackend for FakeChatBackend {
fn tokenizer(&self) -> DynTokenizer {
Arc::new(FakeChatTokenizer)
Arc::clone(&self.tokenizer)
}
fn model_id(&self) -> &str {
@@ -282,65 +265,6 @@ impl ChatRenderer for FakeChatBackend {
}
}
#[derive(Clone, Debug)]
struct FailingDecodeBackend {
inner: FakeChatBackend,
}
#[derive(Debug)]
struct FailingDecodeTokenizer;
impl Tokenizer for FailingDecodeTokenizer {
fn encode(&self, text: &str, add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
FakeChatTokenizer.encode(text, add_special_tokens)
}
fn decode(
&self,
token_ids: &[u32],
skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
if token_ids.contains(&(b'i' as u32)) {
return Err(vllm_tokenizer::TokenizerError("decode failed".to_string()));
}
FakeChatTokenizer.decode(token_ids, skip_special_tokens)
}
fn token_to_id(&self, token: &str) -> Option<u32> {
FakeChatTokenizer.token_to_id(token)
}
}
impl TextBackend for FailingDecodeBackend {
fn tokenizer(&self) -> DynTokenizer {
Arc::new(FailingDecodeTokenizer)
}
fn model_id(&self) -> &str {
self.inner.model_id()
}
}
impl ChatBackend for FailingDecodeBackend {
fn chat_renderer(&self) -> DynChatRenderer {
Arc::new(self.clone())
}
fn new_chat_output_processor(
&self,
_request: &mut ChatRequest,
_options: NewChatOutputProcessorOptions<'_>,
) -> vllm_chat::Result<DynChatOutputProcessor> {
Ok(Box::new(DefaultChatOutputProcessor::plain_text_only()))
}
}
impl ChatRenderer for FailingDecodeBackend {
fn render(&self, request: &ChatRequest) -> vllm_chat::Result<RenderedPrompt> {
self.inner.render(request)
}
}
/// Skip `LogprobsDelta` events that carry only token_ids (no logprobs),
/// returning the next semantically interesting event.
async fn next_semantic<S>(stream: &mut S) -> Option<Result<ChatEvent, vllm_chat::Error>>
@@ -738,7 +662,12 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
send_outputs(
push,
EngineCoreOutputs {
outputs: vec![request_output("chat-4", vec![b'i' as u32], None, None)],
outputs: vec![request_output(
"chat-4",
vec![UNKNOWN_DECODE_TOKEN_ID],
None,
None,
)],
..Default::default()
},
)
@@ -747,9 +676,8 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
},
);
let backend: Arc<dyn ChatTextBackend> = Arc::new(FailingDecodeBackend {
inner: FakeChatBackend::new(),
});
let backend: Arc<dyn ChatTextBackend> =
Arc::new(FakeChatBackend::new().with_tokenizer(Arc::new(TestTokenizer::new())));
let chat = connect_chat_llm_with_ipc(
EngineCoreClientConfig::new_single(handshake_address),
&ipc,
@@ -769,7 +697,10 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
match timeout(Duration::from_secs(2), stream.next()).await.unwrap() {
Some(Err(vllm_chat::Error::Text(vllm_text::Error::Tokenizer(message)))) => {
assert_eq!(message, "decode failed");
assert_eq!(
message,
format!("test tokenizer cannot decode unknown token id {UNKNOWN_DECODE_TOKEN_ID}")
);
}
other => panic!("unexpected event after close: {other:?}"),
}
+158 -63
View File
@@ -1,8 +1,8 @@
//! Text-level roundtrip tests for the real chat-template and output-processor pairing.
//! Roundtrip tests for the real chat-template and output-processor pairing.
//!
//! The invariant under test is that a structured assistant message rendered as history can be
//! parsed from the generated assistant completion and then rendered back to the exact same
//! assistant-completion text.
//! assistant completion.
use std::pin::Pin;
use std::sync::Arc;
@@ -18,6 +18,10 @@ use vllm_chat::{
RendererSelection, load_model_backends,
};
use vllm_text::{DecodedTextEvent, Finished, Prompt};
use vllm_tokenizer::Tokenizer;
const TEXT_COMPLETION_CHUNK_CHARS: usize = 7;
const TOKEN_COMPLETION_CHUNK_TOKENS: usize = 1;
/// One model/parser configuration used to run the fixed roundtrip fixtures.
#[derive(Clone)]
@@ -191,14 +195,28 @@ impl RoundtripCase {
sort_json_keys: false,
}
}
/// GPT-OSS Harmony token-id renderer and native Harmony output processor.
fn gpt_oss() -> Self {
Self {
model_id: "openai/gpt-oss-20b",
assistant_stop_suffix: "", // not applicable for token-id cases
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Always { value: true },
json_fmt: compact_json_fmt(),
sort_json_keys: false,
}
}
}
macro_rules! roundtrip_tests {
($($case:ident => [$($fixture:ident),* $(,)?]),+ $(,)?) => {
($($case:ident => [$($(#[$fixture_attr:meta])* $fixture:ident),* $(,)?]),+ $(,)?) => {
paste::paste! {
$(
$(
#[tokio::test]
$(#[$fixture_attr])*
#[file_serial([<hf_ $case>])]
async fn [<roundtrip_ $case _ $fixture>]() -> Result<()> {
[<run_roundtrip_ $fixture>](RoundtripCase::$case()).await
@@ -217,9 +235,9 @@ roundtrip_tests! {
glm47 => [reasoning_and_content, tool_call_mix],
seed_oss => [reasoning_and_content],
step3p5 => [reasoning_and_content],
gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call
kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history
gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call
}
/// Run the fixed reasoning+content fixture for one model/parser case.
@@ -421,10 +439,10 @@ struct RoundtripResult {
parsed_message: AssistantMessage,
/// Assistant-completion suffix cut from rendering the expected assistant as
/// history.
closed_completion: String,
closed_completion: Prompt,
/// Assistant-completion suffix cut after rendering the parsed assistant
/// back as history.
rerendered_closed_completion: String,
rerendered_closed_completion: Prompt,
}
/// Render, parse, and rerender one assistant turn through the production
@@ -436,60 +454,59 @@ async fn run_roundtrip(
assistant: AssistantMessage,
) -> Result<RoundtripResult> {
let renderer = backends.chat_backend.chat_renderer();
let (prompt, closed_completion_text) =
render_closed_completion(renderer.as_ref(), request, &assistant)?;
let completion_body = closed_completion_text
.strip_suffix(case.assistant_stop_suffix)
.with_context(|| {
format!(
"closed assistant completion did not end with {:?}: {:?}",
case.assistant_stop_suffix, closed_completion_text
)
})?;
let rendered = render_closed_completion(renderer.as_ref(), request, &assistant)?;
let parsed_message =
parse_completion(case, backends, request, &prompt, completion_body).await?;
let (_, rerendered_closed_completion) =
render_closed_completion(renderer.as_ref(), request, &parsed_message)?;
let parsed_message = parse_completion(case, backends, request, &rendered).await?;
let rerendered = render_closed_completion(renderer.as_ref(), request, &parsed_message)?;
Ok(RoundtripResult {
parsed_message,
closed_completion: closed_completion_text,
rerendered_closed_completion,
closed_completion: rendered.completion,
rerendered_closed_completion: rerendered.completion,
})
}
/// Rendered prompt/completion artifacts at the renderer boundary.
struct RenderedTurn {
prompt: Prompt,
completion: Prompt,
}
/// Render `history` as a production prompt and `history + assistant` as closed
/// history, then return the production prompt and assistant-completion suffix.
fn render_closed_completion(
renderer: &dyn vllm_chat::ChatRenderer,
base_request: &ChatRequest,
assistant: &AssistantMessage,
) -> Result<(String, String)> {
) -> Result<RenderedTurn> {
let mut prompt_request = base_request.clone();
prompt_request.chat_options.generation_prompt_mode = GenerationPromptMode::StartNewAssistant;
let prompt = render_text(renderer, &prompt_request).context("failed to render prompt")?;
let prompt = renderer.render(&prompt_request).context("failed to render prompt")?.prompt;
let mut full_request = base_request.clone();
full_request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
full_request.messages.push(ChatMessage::from(assistant.clone()));
let full = render_text(renderer, &full_request).context("failed to render full prompt")?;
let full = renderer.render(&full_request).context("failed to render full prompt")?.prompt;
ensure!(
full.starts_with(&prompt),
"full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
);
let completion = full[prompt.len()..].to_string();
let completion = match (&prompt, full) {
(Prompt::Text(prompt), Prompt::Text(full)) => {
ensure!(
full.starts_with(prompt),
"full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
);
Prompt::Text(full[prompt.len()..].to_string())
}
(Prompt::TokenIds(prompt), Prompt::TokenIds(full)) => {
ensure!(
full.starts_with(prompt),
"full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
);
Prompt::TokenIds(full[prompt.len()..].to_vec())
}
(prompt, full) => bail!("prompt kind changed between renders: {prompt:?} vs {full:?}"),
};
Ok((prompt, completion))
}
/// Render one chat request and require a text prompt.
fn render_text(renderer: &dyn vllm_chat::ChatRenderer, request: &ChatRequest) -> Result<String> {
match renderer.render(request)?.prompt {
Prompt::Text(text) => Ok(text),
other => bail!("roundtrip tests expect text prompts, got {other:?}"),
}
Ok(RenderedTurn { prompt, completion })
}
/// Feed one rendered assistant completion body into the real output processor
@@ -498,13 +515,15 @@ async fn parse_completion(
case: &RoundtripCase,
backends: &vllm_chat::LoadedModelBackends,
base_request: &ChatRequest,
prompt: &str,
completion_body: &str,
rendered: &RenderedTurn,
) -> Result<AssistantMessage> {
let tokenizer = backends.text_backend.tokenizer();
let prompt_token_ids = tokenizer
.encode(prompt, base_request.add_special_tokens)
.context("failed to encode rendered prompt")?;
let prompt_token_ids = match &rendered.prompt {
Prompt::Text(prompt) => tokenizer
.encode(prompt, base_request.add_special_tokens)
.context("failed to encode rendered prompt")?,
Prompt::TokenIds(token_ids) => token_ids.clone(),
};
let mut request = base_request.clone();
let processor = backends.chat_backend.new_chat_output_processor(
@@ -515,7 +534,12 @@ async fn parse_completion(
},
)?;
let decoded = decoded_completion_stream(prompt_token_ids, completion_body);
let decoded = decoded_completion_stream(
tokenizer.as_ref(),
prompt_token_ids,
&rendered.completion,
case.assistant_stop_suffix,
)?;
let mut events = processor.process(decoded)?;
while let Some(event) = events.next().await {
@@ -538,16 +562,46 @@ async fn parse_completion(
/// split into small chunks to exercise streaming parser state across marker
/// and JSON boundaries.
fn decoded_completion_stream(
tokenizer: &dyn Tokenizer,
prompt_token_ids: Vec<u32>,
completion_body: &str,
) -> Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>> {
let prompt_token_count = prompt_token_ids.len();
completion: &Prompt,
assistant_stop_suffix: &str,
) -> Result<Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>>> {
let mut events = vec![DecodedTextEvent::Start {
prompt_token_ids: Arc::from(prompt_token_ids.into_boxed_slice()),
prompt_token_ids: Arc::from(prompt_token_ids.clone().into_boxed_slice()),
prompt_logprobs: None,
}];
let chunks = split_by_chars(completion_body, 7);
let chunks = match completion {
Prompt::Text(text) => {
let body = text.strip_suffix(assistant_stop_suffix).with_context(|| {
format!(
"closed assistant completion did not end with {:?}: {:?}",
assistant_stop_suffix, text
)
})?;
split_by_chars(body, TEXT_COMPLETION_CHUNK_CHARS)
.into_iter()
.map(|delta| DecodedCompletionChunk {
delta,
token_ids: Vec::new(), // unused for text-level roundtrip cases
})
.collect()
}
Prompt::TokenIds(token_ids) => {
ensure!(
assistant_stop_suffix.is_empty(),
"token-id roundtrip cases do not support text stop suffixes"
);
incremental_decode_chunks(
tokenizer,
&prompt_token_ids,
token_ids,
TOKEN_COMPLETION_CHUNK_TOKENS,
)?
}
};
if chunks.is_empty() {
events.push({
DecodedTextEvent::TextDelta {
@@ -555,11 +609,7 @@ fn decoded_completion_stream(
token_ids: Vec::new(),
logprobs: None,
finished: Some(Finished {
usage: vllm_llm::TokenUsage {
prompt_token_count: 0,
output_token_count: 0,
cached_token_count: 0,
},
usage: Default::default(),
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
}),
@@ -569,24 +619,26 @@ fn decoded_completion_stream(
let last_index = chunks.len() - 1;
for (index, chunk) in chunks.into_iter().enumerate() {
let finished = (index == last_index).then(|| Finished {
usage: vllm_llm::TokenUsage {
prompt_token_count,
output_token_count: completion_body.chars().count(),
cached_token_count: 0,
},
usage: Default::default(),
finish_reason: FinishReason::stop_eos(),
kv_transfer_params: None,
});
events.push(DecodedTextEvent::TextDelta {
delta: chunk,
token_ids: Vec::new(),
delta: chunk.delta,
token_ids: chunk.token_ids,
logprobs: None,
finished,
});
}
}
stream::iter(events).map(Ok).boxed()
Ok(stream::iter(events).map(Ok).boxed())
}
/// One decoded completion chunk fed into the output processor.
struct DecodedCompletionChunk {
delta: String,
token_ids: Vec<u32>,
}
/// Split text into chunks containing at most `chunk_chars` Unicode scalar
@@ -612,6 +664,49 @@ fn split_by_chars(text: &str, chunk_chars: usize) -> Vec<String> {
chunks
}
/// Split token ids into chunks containing at most `chunk_size` ids.
fn split_by_count(token_ids: &[u32], chunk_size: usize) -> Vec<Vec<u32>> {
token_ids.chunks(chunk_size).map(<[u32]>::to_vec).collect()
}
/// Decode token ids incrementally using the production tokenizer stream.
fn incremental_decode_chunks(
tokenizer: &dyn Tokenizer,
prompt_token_ids: &[u32],
token_ids: &[u32],
chunk_size: usize,
) -> Result<Vec<DecodedCompletionChunk>> {
let mut decoder = tokenizer.create_decode_stream(prompt_token_ids, false, 0);
let mut chunks = Vec::new();
for chunk_token_ids in split_by_count(token_ids, chunk_size) {
let mut delta = String::new();
for token_id in chunk_token_ids.iter().copied() {
decoder.push_token(token_id)?;
while let Some(chunk) = decoder.next_chunk() {
delta.push_str(&chunk);
}
}
chunks.push(DecodedCompletionChunk {
delta,
token_ids: chunk_token_ids,
});
}
let (last_chunk, _) = decoder.flush(None)?;
if let Some(last_chunk) = last_chunk {
if let Some(delta) = chunks.last_mut() {
delta.delta.push_str(&last_chunk);
} else {
chunks.push(DecodedCompletionChunk {
delta: last_chunk,
token_ids: Vec::new(),
});
}
}
Ok(chunks)
}
/// Build a chat request fixture with parser-enabling tool-choice semantics.
fn roundtrip_request(
request_id: impl Into<String>,
+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);
}
}
}
+1
View File
@@ -23,6 +23,7 @@ expect-test.workspace = true
futures.workspace = true
openai-protocol.workspace = true
tool-parser.workspace = true
vllm-tokenizer = { workspace = true, features = ["test-utils"] }
[[bench]]
name = "deepseek_v3"
+4
View File
@@ -27,6 +27,10 @@ impl Tokenizer for BenchTokenizer {
fn token_to_id(&self, _token: &str) -> Option<u32> {
Some(u32::MAX)
}
fn id_to_token(&self, _id: u32) -> Option<String> {
Some("\u{FFFD}".to_string())
}
}
/// Bench-only adapter that exposes a unified parser through the tool-parser
+15 -12
View File
@@ -49,11 +49,14 @@ mod tests {
use std::sync::Arc;
use super::SeedOssReasoningParser;
use crate::reasoning::{ReasoningParser, tests::FakeTokenizer};
use crate::reasoning::{
ReasoningParser,
tests::{SEED_THINK_END_ID, SEED_THINK_START_ID, fake_tokenizer},
};
#[test]
fn without_prompt_markers_expects_start_token() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("implicit reasoning</seed:think>answer").unwrap();
@@ -66,10 +69,10 @@ mod tests {
#[test]
fn picks_up_prompt_start_boundary() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
// Prompt prefills `<seed:think>` (id 10), opening reasoning before the stream.
parser.initialize(&[10]).unwrap();
// Prompt prefills `<seed:think>`, opening reasoning before the stream.
parser.initialize(&[SEED_THINK_START_ID]).unwrap();
let delta = parser.push("reason</seed:think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
@@ -78,10 +81,10 @@ mod tests {
#[test]
fn respects_prompt_end_boundary() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
// Prompt already closed reasoning with `</seed:think>` (id 11).
parser.initialize(&[11]).unwrap();
// Prompt already closed reasoning with `</seed:think>`.
parser.initialize(&[SEED_THINK_END_ID]).unwrap();
let delta = parser.push("answer").unwrap();
assert_eq!(delta.reasoning, None);
@@ -91,7 +94,7 @@ mod tests {
#[test]
fn handles_explicit_start_token() {
// An explicit start delimiter must not leak into reasoning text.
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<seed:think>reason</seed:think>answer").unwrap();
@@ -103,7 +106,7 @@ mod tests {
fn streams_explicit_start_token_across_pushes() {
// Start token, reasoning body, end token, and content arrive in separate
// streaming deltas.
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
let mut reasoning = String::new();
@@ -131,9 +134,9 @@ mod tests {
#[test]
fn handles_partial_delimiters_across_pushes() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[10]).unwrap();
parser.initialize(&[SEED_THINK_START_ID]).unwrap();
// Closing delimiter `</seed:think>` arrives in two halves.
let first = parser.push("reason</seed:").unwrap();
+20 -17
View File
@@ -127,14 +127,17 @@ mod tests {
use std::sync::Arc;
use super::Step3p5ReasoningParser;
use crate::reasoning::{ReasoningParser, tests::FakeTokenizer};
use crate::reasoning::{
ReasoningParser,
tests::{THINK_START_ID, fake_tokenizer},
};
#[test]
fn picks_up_prompt_start_boundary() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
// Prompt prefills `<think>` (id 1), opening reasoning before the stream.
parser.initialize(&[1]).unwrap();
// Prompt prefills `<think>`, opening reasoning before the stream.
parser.initialize(&[THINK_START_ID]).unwrap();
let delta = parser.push("This is a reasoning section</think>This is the rest").unwrap();
assert_eq!(
@@ -146,7 +149,7 @@ mod tests {
#[test]
fn handles_unterminated_reasoning() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let pushed = parser.push("<think>reason without end").unwrap();
@@ -159,7 +162,7 @@ mod tests {
#[test]
fn handles_empty_input() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let pushed = parser.push("").unwrap();
@@ -172,9 +175,9 @@ mod tests {
fn complex_newline_pattern_trims_only_single_framing_newline_each_side() {
// Only the immediately-adjacent framing `\n` is dropped on each side of
// `</think>`; surrounding newlines remain part of reasoning/content.
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[1]).unwrap();
parser.initialize(&[THINK_START_ID]).unwrap();
let delta = parser
.push("\n This is a \n reasoning section\n\n\n</think>\n\nThis is the rest")
@@ -188,7 +191,7 @@ mod tests {
#[test]
fn drops_framing_newlines_in_single_push() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>reason\n</think>\nanswer").unwrap();
@@ -198,7 +201,7 @@ mod tests {
#[test]
fn drops_framing_newlines_across_pushes() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
// The trailing `\n` from the first push is held until we know whether
@@ -219,7 +222,7 @@ mod tests {
#[test]
fn replays_held_newline_when_more_reasoning_follows() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let first = parser.push("<think>reason\n").unwrap();
@@ -232,7 +235,7 @@ mod tests {
#[test]
fn finish_flushes_held_newline_in_unterminated_stream() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let first = parser.push("<think>reason\n").unwrap();
@@ -245,7 +248,7 @@ mod tests {
#[test]
fn preserves_inner_newlines_in_reasoning() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>line1\nline2</think>tail").unwrap();
@@ -257,7 +260,7 @@ mod tests {
fn trims_only_one_trailing_reasoning_newline() {
// Only the single framing newline immediately before `</think>` is
// dropped; earlier newlines in the reasoning body are preserved.
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>reason\n\n</think>answer").unwrap();
@@ -269,7 +272,7 @@ mod tests {
fn drops_only_first_content_newline_after_transition() {
// The leading-`\n` drop applies only to the first content delta after
// `</think>`; later deltas pass through untouched.
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let first = parser.push("<think>reason</think>").unwrap();
@@ -288,7 +291,7 @@ mod tests {
#[test]
fn passes_through_clean_boundary_without_framing_newlines() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>reason</think>tail").unwrap();
@@ -298,7 +301,7 @@ mod tests {
#[test]
fn handles_empty_reasoning_section() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think></think>answer").unwrap();
+48 -60
View File
@@ -1,54 +1,42 @@
use std::sync::Arc;
use vllm_tokenizer::Tokenizer;
use vllm_tokenizer::test_utils::TestTokenizer;
use super::{
DeepSeekR1ReasoningParser, DelimitedReasoningParser, MiniMaxM3ReasoningParser,
Qwen3ReasoningParser, ReasoningParser,
};
pub(crate) struct FakeTokenizer;
pub(crate) const THINK_START_ID: u32 = 256;
pub(crate) const THINK_END_ID: u32 = 257;
pub(crate) const START_THINKING_ID: u32 = 258;
pub(crate) const END_THINKING_ID: u32 = 259;
pub(crate) const MINIMAX_THINK_START_ID: u32 = 260;
pub(crate) const MINIMAX_THINK_END_ID: u32 = 261;
pub(crate) const SPECIAL_BOUNDARY_ID: u32 = 262;
pub(crate) const MM_THINK_START_ID: u32 = 263;
pub(crate) const MM_THINK_END_ID: u32 = 264;
pub(crate) const SEED_THINK_START_ID: u32 = 265;
pub(crate) const SEED_THINK_END_ID: u32 = 266;
impl Tokenizer for FakeTokenizer {
fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok(token_ids
.iter()
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
.collect())
}
fn token_to_id(&self, token: &str) -> Option<u32> {
match token {
"<think>" => Some(1),
"</think>" => Some(2),
"<|START_THINKING|>" => Some(3),
"<|END_THINKING|>" => Some(4),
"◁think▷" => Some(5),
"◁/think▷" => Some(6),
"<mm:think>" => Some(8),
"</mm:think>" => Some(9),
"<seed:think>" => Some(10),
"</seed:think>" => Some(11),
_ => None,
}
}
fn is_special_id(&self, token_id: u32) -> bool {
token_id == 7
}
pub(crate) fn fake_tokenizer() -> TestTokenizer {
TestTokenizer::new()
.with_regular_token("<think>", THINK_START_ID)
.with_regular_token("</think>", THINK_END_ID)
.with_regular_token("<|START_THINKING|>", START_THINKING_ID)
.with_regular_token("<|END_THINKING|>", END_THINKING_ID)
.with_regular_token("◁think▷", MINIMAX_THINK_START_ID)
.with_regular_token("◁/think▷", MINIMAX_THINK_END_ID)
.with_special_token("<special-boundary>", SPECIAL_BOUNDARY_ID)
.with_regular_token("<mm:think>", MM_THINK_START_ID)
.with_regular_token("</mm:think>", MM_THINK_END_ID)
.with_regular_token("<seed:think>", SEED_THINK_START_ID)
.with_regular_token("</seed:think>", SEED_THINK_END_ID)
}
#[test]
fn delimited_content_only_stream() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser =
DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
@@ -60,7 +48,7 @@ fn delimited_content_only_stream() {
#[test]
fn delimited_single_chunk_with_reasoning_and_content() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser =
DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
@@ -71,7 +59,7 @@ fn delimited_single_chunk_with_reasoning_and_content() {
#[test]
fn delimited_partial_tokens_across_chunks() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser =
DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
@@ -83,10 +71,10 @@ fn delimited_partial_tokens_across_chunks() {
#[test]
fn delimited_finish_flushes_buffer() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser =
DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
parser.initialize(&[1]);
parser.initialize(&[THINK_START_ID]);
let delta = parser.push("unfinished</thi");
assert_eq!(delta.reasoning.as_deref(), Some("unfinished"));
@@ -96,7 +84,7 @@ fn delimited_finish_flushes_buffer() {
#[test]
fn qwen3_without_prompt_markers_expects_start_token() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("reason</think>answer").unwrap();
@@ -106,9 +94,9 @@ fn qwen3_without_prompt_markers_expects_start_token() {
#[test]
fn qwen3_prompt_end_marker_starts_in_content() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[2]).unwrap();
parser.initialize(&[THINK_END_ID]).unwrap();
let delta = parser.push("answer").unwrap();
assert_eq!(delta.reasoning, None);
@@ -117,7 +105,7 @@ fn qwen3_prompt_end_marker_starts_in_content() {
#[test]
fn qwen3_tolerates_old_and_new_formats() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut old_parser = Qwen3ReasoningParser::new(tokenizer.clone()).unwrap();
let old = old_parser.push("<think>reason</think>answer").unwrap();
@@ -125,7 +113,7 @@ fn qwen3_tolerates_old_and_new_formats() {
assert_eq!(old.content.as_deref(), Some("answer"));
let mut new_parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
new_parser.initialize(&[1]).unwrap();
new_parser.initialize(&[THINK_START_ID]).unwrap();
let new = new_parser.push("reason</think>answer").unwrap();
assert_eq!(new.reasoning.as_deref(), Some("reason"));
assert_eq!(new.content.as_deref(), Some("answer"));
@@ -133,10 +121,10 @@ fn qwen3_tolerates_old_and_new_formats() {
#[test]
fn qwen3_stops_scanning_at_last_special_token() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[1, 7]).unwrap();
parser.initialize(&[THINK_START_ID, SPECIAL_BOUNDARY_ID]).unwrap();
let delta = parser.push("answer").unwrap();
assert_eq!(delta.reasoning, None);
@@ -145,7 +133,7 @@ fn qwen3_stops_scanning_at_last_special_token() {
#[test]
fn deepseek_r1_defaults_to_reasoning_without_prompt_boundary() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = DeepSeekR1ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("reason</think>answer").unwrap();
@@ -155,10 +143,10 @@ fn deepseek_r1_defaults_to_reasoning_without_prompt_boundary() {
#[test]
fn deepseek_r1_stops_scanning_at_last_special_token() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = DeepSeekR1ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[2, 7]).unwrap();
parser.initialize(&[THINK_END_ID, SPECIAL_BOUNDARY_ID]).unwrap();
let delta = parser.push("reason</think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
@@ -167,7 +155,7 @@ fn deepseek_r1_stops_scanning_at_last_special_token() {
#[test]
fn minimax_m3_handles_explicit_think_delimiters() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<mm:think>reason</mm:think>answer").unwrap();
@@ -177,7 +165,7 @@ fn minimax_m3_handles_explicit_think_delimiters() {
#[test]
fn minimax_m3_drops_leading_end_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("</mm:think>answer").unwrap();
@@ -187,7 +175,7 @@ fn minimax_m3_drops_leading_end_marker() {
#[test]
fn minimax_m3_preserves_non_leading_end_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("XXX</mm:think>YYY").unwrap();
@@ -197,7 +185,7 @@ fn minimax_m3_preserves_non_leading_end_marker() {
#[test]
fn minimax_m3_drops_split_leading_end_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
assert!(parser.push("</mm").unwrap().is_empty());
@@ -208,9 +196,9 @@ fn minimax_m3_drops_split_leading_end_marker() {
#[test]
fn minimax_m3_uses_prompt_prefilled_start_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[8]).unwrap();
parser.initialize(&[MM_THINK_START_ID]).unwrap();
let delta = parser.push("reason</mm:think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
@@ -219,9 +207,9 @@ fn minimax_m3_uses_prompt_prefilled_start_marker() {
#[test]
fn minimax_m3_uses_prompt_prefilled_end_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(fake_tokenizer());
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[9]).unwrap();
parser.initialize(&[MM_THINK_END_ID]).unwrap();
let delta = parser.push("answer").unwrap();
assert_eq!(delta.reasoning, 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]
+6 -31
View File
@@ -124,42 +124,17 @@ impl UnifiedParser for CombinedParser {
mod tests {
use std::sync::Arc;
use vllm_tokenizer::Tokenizer;
use vllm_tokenizer::test_utils::TestTokenizer;
use super::CombinedParser;
use crate::reasoning::{Qwen3ReasoningParser, ReasoningDelta, ReasoningParser};
use crate::tool::{Qwen3XmlToolParser, Tool, ToolParser};
use crate::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};
struct FakeTokenizer;
impl Tokenizer for FakeTokenizer {
fn encode(
&self,
text: &str,
_add_special_tokens: bool,
) -> vllm_tokenizer::Result<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok(token_ids
.iter()
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
.collect())
}
fn token_to_id(&self, token: &str) -> Option<u32> {
match token {
"<think>" => Some(1),
"</think>" => Some(2),
_ => None,
}
}
fn tokenizer() -> TestTokenizer {
TestTokenizer::new()
.with_regular_token("<think>", 256)
.with_regular_token("</think>", 257)
}
fn test_tools() -> Vec<Tool> {
@@ -273,7 +248,7 @@ mod tests {
#[test]
fn combined_parser_emits_reasoning_and_text() {
let tokenizer = Arc::new(FakeTokenizer);
let tokenizer = Arc::new(tokenizer());
let reasoning = Qwen3ReasoningParser::create(tokenizer).unwrap();
let mut parser = CombinedParser::new(Some(reasoning), None);
+15 -66
View File
@@ -515,7 +515,7 @@ mod tests {
use serde_json::{Value, json};
use thiserror_ext::AsReport;
use vllm_tokenizer::Tokenizer;
use vllm_tokenizer::test_utils::TestTokenizer;
use winnow::combinator::{eof, terminated};
use winnow::error::ErrMode;
use winnow::prelude::*;
@@ -527,66 +527,15 @@ mod tests {
use crate::tool::Tool;
use crate::unified::{UnifiedParserEvent, parsing_failed};
struct FakeTokenizer;
const CHANNEL_START_ID: u32 = 256;
const CHANNEL_END_ID: u32 = 257;
const TURN_BOUNDARY_ID: u32 = 258;
impl Tokenizer for FakeTokenizer {
fn encode(
&self,
text: &str,
_add_special_tokens: bool,
) -> vllm_tokenizer::Result<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok(token_ids
.iter()
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
.collect())
}
fn token_to_id(&self, token: &str) -> Option<u32> {
match token {
CHANNEL_START => Some(100),
CHANNEL_END => Some(101),
_ => None,
}
}
fn is_special_id(&self, token_id: u32) -> bool {
matches!(token_id, 100..=105)
}
}
struct MissingTokenTokenizer;
impl Tokenizer for MissingTokenTokenizer {
fn encode(
&self,
text: &str,
_add_special_tokens: bool,
) -> vllm_tokenizer::Result<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok(token_ids
.iter()
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
.collect())
}
fn token_to_id(&self, _token: &str) -> Option<u32> {
None
}
fn tokenizer() -> TestTokenizer {
TestTokenizer::new()
.with_special_token(CHANNEL_START, CHANNEL_START_ID)
.with_special_token(CHANNEL_END, CHANNEL_END_ID)
.with_special_token("<turn-boundary>", TURN_BOUNDARY_ID)
}
trait UnifiedParserTestExt {
@@ -716,12 +665,12 @@ mod tests {
}
fn test_parser() -> Gemma4UnifiedParser {
Gemma4UnifiedParser::new(&test_tools(), Arc::new(FakeTokenizer)).unwrap()
Gemma4UnifiedParser::new(&test_tools(), Arc::new(tokenizer())).unwrap()
}
#[test]
fn gemma4_create_requires_channel_start_token() {
let error = match Gemma4UnifiedParser::new(&test_tools(), Arc::new(MissingTokenTokenizer)) {
let error = match Gemma4UnifiedParser::new(&test_tools(), Arc::new(TestTokenizer::new())) {
Ok(_) => panic!("expected missing token error"),
Err(error) => error,
};
@@ -1046,7 +995,7 @@ mod tests {
#[test]
fn gemma4_initialize_open_channel_prompt_starts_in_reasoning() {
let mut parser = test_parser();
parser.initialize(&[100, 3000, 3001]).unwrap();
parser.initialize(&[CHANNEL_START_ID, 3000, 3001]).unwrap();
let output = parser.parse_complete("reason<channel|>answer").unwrap();
@@ -1057,7 +1006,7 @@ mod tests {
#[test]
fn gemma4_initialize_turn_prompt_starts_in_text() {
let mut parser = test_parser();
parser.initialize(&[104, 3000, 3001]).unwrap();
parser.initialize(&[TURN_BOUNDARY_ID, 3000, 3001]).unwrap();
let output = parser.parse_complete("<|channel>thought\nreason<channel|>answer").unwrap();
@@ -1068,7 +1017,7 @@ mod tests {
#[test]
fn gemma4_initialize_special_token_caps_boundary_scan() {
let mut parser = test_parser();
parser.initialize(&[100, 3000, 104, 3001]).unwrap();
parser.initialize(&[CHANNEL_START_ID, 3000, TURN_BOUNDARY_ID, 3001]).unwrap();
let output = parser.parse_complete("answer").unwrap();
@@ -1079,7 +1028,7 @@ mod tests {
#[test]
fn gemma4_initialize_closed_channel_prompt_starts_in_text() {
let mut parser = test_parser();
parser.initialize(&[100, 3000, 3001, 101]).unwrap();
parser.initialize(&[CHANNEL_START_ID, 3000, 3001, CHANNEL_END_ID]).unwrap();
let output = parser.parse_complete("answer").unwrap();
+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"
);
}
}
+2
View File
@@ -61,8 +61,10 @@ expect-test.workspace = true
rmp-serde.workspace = true
serial_test.workspace = true
tempfile.workspace = true
tokio = { workspace = true, features = ["test-util"] }
tower.workspace = true
vllm-engine-core-client = { workspace = true, features = ["test-util"] }
vllm-tokenizer = { workspace = true, features = ["test-utils"] }
zeromq.workspace = true
[lints]
+1 -78
View File
@@ -4,21 +4,16 @@ mod convert;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use futures::{Stream, StreamExt as _, stream};
use futures::{Stream, StreamExt as _};
use thiserror_ext::AsReport as _;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::sync::mpsc;
use tokio_openssl::SslStream;
use tokio_stream::wrappers::ReceiverStream;
use tonic::transport::server::{Connected, TcpConnectInfo};
use tonic::{Request, Response, Status};
use tracing::info;
use vllm_text::{DecodedTextEvent, TextOutputStreamExt as _};
use self::convert::ResponseOpts;
use crate::listener::{Listener, ListenerIo};
use crate::state::AppState;
/// Generated protobuf/gRPC types for the `vllm` package.
@@ -31,78 +26,6 @@ pub use pb::generate_server::GenerateServer;
#[cfg(test)]
mod tests;
/// Newtype over `tokio-openssl`'s `SslStream` so we can implement tonic's
/// [`Connected`] on it (the orphan rule blocks doing so on the foreign type).
pub(crate) struct GrpcTlsStream {
inner: SslStream<ListenerIo>,
}
impl GrpcTlsStream {
pub(crate) fn new(inner: SslStream<ListenerIo>) -> Self {
Self { inner }
}
}
impl AsyncRead for GrpcTlsStream {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
}
}
impl AsyncWrite for GrpcTlsStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.get_mut().inner).poll_flush(cx)
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
}
}
impl Connected for GrpcTlsStream {
type ConnectInfo = TcpConnectInfo;
fn connect_info(&self) -> TcpConnectInfo {
self.inner.get_ref().connect_info()
}
}
/// Adapt the shared server listener into tonic's incoming stream shape.
pub(crate) fn incoming(listener: Listener) -> impl Stream<Item = std::io::Result<ListenerIo>> {
stream::unfold(listener, |mut listener| async move {
let (io, _) = axum::serve::Listener::accept(&mut listener).await;
Some((Ok(io), listener))
})
}
/// Wrap the gRPC listener so each accepted connection completes a TLS handshake
/// before tonic serves it.
pub(crate) fn tls_incoming(
listener: Listener,
context: openssl::ssl::SslContext,
handshake_timeout: std::time::Duration,
) -> impl Stream<Item = std::io::Result<GrpcTlsStream>> {
tls_listener::builder(context)
.handshake_timeout(handshake_timeout)
.listen(listener)
.map(|res| {
res.map(|(inner, _addr)| GrpcTlsStream::new(inner))
.map_err(std::io::Error::other)
})
}
/// gRPC Generate service implementation backed by the shared application state.
pub struct GenerateServiceImpl {
state: Arc<AppState>,
+8 -35
View File
@@ -24,14 +24,15 @@ use vllm_engine_core_client::protocol::{
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
use vllm_llm::Llm;
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
use vllm_text::tokenizer::DynTokenizer;
use vllm_text::{Prompt, TextBackend};
use vllm_tokenizer::test_utils::TestTokenizer;
use zeromq::prelude::{SocketRecv, SocketSend};
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
use super::pb::generate_client::GenerateClient;
use super::{GenerateServer, GenerateServiceImpl, incoming, pb, tls_incoming};
use crate::listener::Listener;
use super::{GenerateServer, GenerateServiceImpl, pb};
use crate::listener::{Listener, MaybeTlsListener};
use crate::state::AppState;
use crate::tls;
use crate::tls_tests::{TestCerts, server_tls};
@@ -155,37 +156,9 @@ fn test_llm(client: EngineCoreClient) -> Llm {
#[derive(Clone, Debug)]
struct FakeTextBackend;
#[derive(Debug)]
struct FakeTokenizer;
impl Tokenizer for FakeTokenizer {
fn encode(
&self,
text: &str,
_add_special_tokens: bool,
) -> vllm_text::tokenizer::Result<Vec<u32>> {
Ok(text.bytes().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_text::tokenizer::Result<String> {
Ok(
String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
.into_owned(),
)
}
fn token_to_id(&self, token: &str) -> Option<u32> {
token.bytes().next().map(u32::from)
}
}
impl TextBackend for FakeTextBackend {
fn tokenizer(&self) -> DynTokenizer {
Arc::new(FakeTokenizer)
Arc::new(TestTokenizer::new())
}
fn model_id(&self) -> &str {
@@ -287,7 +260,7 @@ async fn grpc_test_server(
let addr = listener.local_addr().expect("local addr");
let server_task = tokio::spawn(async move {
let incoming = incoming(Listener::Tcp(listener));
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
TonicServer::builder()
.add_service(svc)
.serve_with_incoming(incoming)
@@ -318,7 +291,7 @@ async fn grpc_tls_test_server(
let addr = listener.local_addr().expect("local addr").to_string();
let server_task = tokio::spawn(async move {
let incoming = tls_incoming(Listener::Tcp(listener), context, tls::TLS_HANDSHAKE_TIMEOUT);
let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context);
TonicServer::builder()
.add_service(svc)
.serve_with_incoming(incoming)
@@ -413,7 +386,7 @@ async fn grpc_server_with_keepalive(
}
let server_task = tokio::spawn(async move {
let incoming = incoming(Listener::Tcp(listener));
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
builder
.add_service(svc)
.serve_with_incoming(incoming)
+17 -62
View File
@@ -27,7 +27,6 @@ pub use config::{
ApiServerOptions, Config, CoordinatorMode, CorsConfig, DEFAULT_KEEP_ALIVE_TIMEOUT,
HttpListenerMode, TlsConfig,
};
use futures::FutureExt as _;
use hyper::body::Incoming;
use hyper::server::conn::http1;
use hyper_util::rt::{TokioIo, TokioTimer};
@@ -45,7 +44,7 @@ use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
use vllm_llm::Llm;
use vllm_text::TextLlm;
use crate::listener::Listener;
use crate::listener::{Listener, MaybeTlsListener};
use crate::routes::build_router;
use crate::server_info::ServerInfoSnapshot;
use crate::state::AppState;
@@ -179,7 +178,7 @@ where
let listener = Listener::bind(&config.listener_mode)
.await
.context("failed to bind listener for OpenAI server")?;
let bind_address = listener.local_addr()?;
let bind_address = listener.local_addr_display()?;
let model = state.primary_model_name().to_owned();
let app = extend_router(build_router(state.clone()));
@@ -252,7 +251,6 @@ where
// silent client cannot hold the connection open.
let keep_alive_timeout = config.keep_alive_timeout;
let timeouts = ConnectionTimeouts {
handshake: tls::TLS_HANDSHAKE_TIMEOUT,
header_read: if keep_alive_timeout.is_zero() {
DEFAULT_KEEP_ALIVE_TIMEOUT
} else {
@@ -266,9 +264,15 @@ where
let server_shutdown = server_shutdown.clone();
let force_shutdown = force_shutdown.clone();
async move {
let listener = match tls_config {
Some(context) => MaybeTlsListener::tls(listener, context),
None => MaybeTlsListener::plain(listener),
};
let server = serve_connections(listener, app, shutdown.cancelled_owned(), timeouts);
let result = tokio::select! {
result = serve_listener(listener, tls_config, app, shutdown.cancelled_owned(), timeouts) => {
result
result = server => {
result.context("HTTP server failed")
}
_ = force_shutdown.cancelled() => {
warn!("HTTP graceful shutdown deadline elapsed; aborting server");
@@ -292,18 +296,11 @@ where
shutdown.cancelled().await;
return Ok(());
};
// Box to unify the TLS and plaintext arms' different stream types.
let server = match grpc_tls {
Some(context) => {
let incoming =
grpc::tls_incoming(grpc_listener, context, tls::TLS_HANDSHAKE_TIMEOUT);
svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned()).boxed()
}
None => {
let incoming = grpc::incoming(grpc_listener);
svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned()).boxed()
}
let incoming = match grpc_tls {
Some(context) => MaybeTlsListener::tls(grpc_listener, context),
None => MaybeTlsListener::plain(grpc_listener),
};
let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());
let result = tokio::select! {
result = server => {
@@ -333,61 +330,19 @@ where
/// Per-connection timeouts applied while serving HTTP/HTTPS.
#[derive(Clone, Copy)]
pub(crate) struct ConnectionTimeouts {
/// Max time for a client to complete the TLS handshake (TLS path only).
pub(crate) handshake: Duration,
/// HTTP/1 header-read timeout (bounds idle keep-alive and the head read).
pub(crate) header_read: Duration,
/// Whether HTTP/1 keep-alive is enabled; `false` closes after each response.
pub(crate) keep_alive_enabled: bool,
}
/// Apply optional TLS termination and per-connection HTTP timeouts, then serve
/// `app`. Shared by [`serve_with_router_extension`] and the TLS tests.
async fn serve_listener(
listener: Listener,
tls: Option<openssl::ssl::SslContext>,
app: Router,
shutdown: impl Future<Output = ()> + Send + 'static,
timeouts: ConnectionTimeouts,
) -> Result<()> {
match tls {
Some(context) => {
// tls-listener terminates TLS (handshake + timeout); serve_connections
// owns the HTTP keep-alive/idle bound that axum::serve cannot express.
// Failed handshakes (incl. timeouts) log at ERROR via tls-listener.
let listener = tls_listener::builder(context)
.handshake_timeout(timeouts.handshake)
.listen(listener);
serve_connections(
listener,
app,
shutdown,
timeouts.header_read,
timeouts.keep_alive_enabled,
)
.await
.context("HTTPS server failed")
}
None => serve_connections(
listener,
app,
shutdown,
timeouts.header_read,
timeouts.keep_alive_enabled,
)
.await
.context("HTTP server failed"),
}
}
/// Serve `app` per connection (HTTP/1) with a keep-alive idle timeout and
/// graceful drain. Hand-rolled on hyper because [`axum::serve()`] takes no config.
async fn serve_connections<L>(
mut listener: L,
app: Router,
shutdown: impl Future<Output = ()> + Send,
header_read: Duration,
keep_alive_enabled: bool,
timeouts: ConnectionTimeouts,
) -> Result<()>
where
L: axum::serve::Listener,
@@ -404,8 +359,8 @@ where
app.clone().map_request(|req: Request<Incoming>| req.map(Body::new)),
);
let mut builder = http1::Builder::new();
builder.timer(TokioTimer::new()).header_read_timeout(header_read);
if !keep_alive_enabled {
builder.timer(TokioTimer::new()).header_read_timeout(timeouts.header_read);
if !timeouts.keep_alive_enabled {
builder.keep_alive(false);
}
let connection = builder.serve_connection(TokioIo::new(io), service);
+96 -6
View File
@@ -12,13 +12,14 @@ use std::pin::Pin;
use std::task::{Context, Poll, ready};
use auto_enums::enum_derive;
use openssl::ssl::SslContext;
use socket2::Socket;
use tls_listener::{AsyncAccept, AsyncListener};
use tokio::net::{TcpListener, TcpStream, UnixListener, UnixStream};
use tonic::transport::server::{Connected, TcpConnectInfo};
use tracing::trace;
use crate::HttpListenerMode;
use crate::{HttpListenerMode, tls};
/// Runtime listener type used by the OpenAI-compatible HTTP or gRPC server,
/// which is either a TCP listener or a Unix-domain listener.
@@ -61,7 +62,7 @@ impl Listener {
/// Return a log-friendly local address string for either TCP or Unix
/// sockets.
pub fn local_addr(&self) -> Result<String> {
pub fn local_addr_display(&self) -> Result<String> {
match self {
Self::Tcp(listener) => Ok(listener.local_addr()?.to_string()),
Self::Unix(listener) => Ok(match listener.local_addr()?.as_pathname() {
@@ -92,7 +93,7 @@ impl Listener {
}
}
fn listener_addr(&self) -> Result<ListenerAddr> {
fn local_addr(&self) -> Result<ListenerAddr> {
match self {
Self::Tcp(listener) => listener.local_addr().map(ListenerAddr::Tcp),
Self::Unix(listener) => listener.local_addr().map(ListenerAddr::Unix),
@@ -100,6 +101,7 @@ impl Listener {
}
}
/// Allow the unified listener to plug directly into tonic's gRPC server.
impl Connected for ListenerIo {
type ConnectInfo = TcpConnectInfo;
@@ -144,7 +146,7 @@ impl axum::serve::Listener for Listener {
}
fn local_addr(&self) -> Result<Self::Addr> {
self.listener_addr()
self.local_addr()
}
}
@@ -173,10 +175,98 @@ impl AsyncAccept for Listener {
}
}
}
impl AsyncListener for Listener {
fn local_addr(&self) -> Result<Self::Address> {
self.listener_addr()
self.local_addr()
}
}
/// A listener that may be either a plain TCP/UDS listener or a TLS listener over it.
pub enum MaybeTlsListener {
Plain(Listener),
Tls(tls_listener::TlsListener<Listener, SslContext>),
}
impl MaybeTlsListener {
/// Create a plain listener without TLS.
pub fn plain(listener: Listener) -> Self {
Self::Plain(listener)
}
/// Create a TLS listener over the given plain listener.
pub fn tls(listener: Listener, context: SslContext) -> Self {
Self::Tls(
tls_listener::builder(context)
.handshake_timeout(tls::TLS_HANDSHAKE_TIMEOUT)
.listen(listener),
)
}
}
/// Listener I/O type that may be either a plain TCP/UDS stream or a TLS stream over it.
#[derive(Debug)]
#[enum_derive(tokio1::AsyncRead, tokio1::AsyncWrite)]
pub enum MaybeTlsStream {
Plain(ListenerIo),
Tls(tokio_openssl::SslStream<ListenerIo>),
}
/// Allow the maybe-TLS listener to plug directly into `axum::serve(...)`.
impl axum::serve::Listener for MaybeTlsListener {
type Addr = ListenerAddr;
type Io = MaybeTlsStream;
async fn accept(&mut self) -> (Self::Io, Self::Addr) {
match self {
Self::Plain(listener) => {
let (io, addr) = axum::serve::Listener::accept(listener).await;
(MaybeTlsStream::Plain(io), addr)
}
Self::Tls(tls_listener) => {
let (io, addr) = axum::serve::Listener::accept(tls_listener).await;
(MaybeTlsStream::Tls(io), addr)
}
}
}
fn local_addr(&self) -> tokio::io::Result<Self::Addr> {
match self {
Self::Plain(listener) => listener.local_addr(),
Self::Tls(tls_listener) => tls_listener.local_addr(),
}
}
}
/// Allow the maybe-TLS listener to plug directly into tonic's gRPC server.
impl Connected for MaybeTlsStream {
type ConnectInfo = TcpConnectInfo;
fn connect_info(&self) -> TcpConnectInfo {
match self {
Self::Plain(stream) => stream.connect_info(),
Self::Tls(stream) => stream.get_ref().connect_info(),
}
}
}
/// Allow the maybe-TLS listener to be adaptable to tonic's incoming stream shape.
impl futures::Stream for MaybeTlsListener {
type Item = std::io::Result<MaybeTlsStream>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
match self.get_mut() {
Self::Plain(listener) => {
let listener = Pin::new(listener);
let (io, _) = ready!(listener.poll_accept(cx))?;
Poll::Ready(Some(Ok(MaybeTlsStream::Plain(io))))
}
Self::Tls(tls_listener) => {
let tls_listener = Pin::new(tls_listener);
let (io, _) =
ready!(tls_listener.poll_accept(cx)).map_err(std::io::Error::other)?;
Poll::Ready(Some(Ok(MaybeTlsStream::Tls(io))))
}
}
}
}
@@ -24,8 +24,9 @@ use vllm_engine_core_client::protocol::{
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
use vllm_llm::Llm;
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
use vllm_text::tokenizer::DynTokenizer;
use vllm_text::{Prompt, TextBackend};
use vllm_tokenizer::test_utils::TestTokenizer;
use zeromq::prelude::{SocketRecv, SocketSend};
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
@@ -151,37 +152,9 @@ fn test_llm(client: EngineCoreClient) -> Llm {
#[derive(Clone, Debug)]
struct FakeChatBackend;
#[derive(Debug)]
struct FakeChatTokenizer;
impl Tokenizer for FakeChatTokenizer {
fn encode(
&self,
text: &str,
_add_special_tokens: bool,
) -> vllm_text::tokenizer::Result<Vec<u32>> {
Ok(text.bytes().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_text::tokenizer::Result<String> {
Ok(
String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
.into_owned(),
)
}
fn token_to_id(&self, token: &str) -> Option<u32> {
token.bytes().next().map(u32::from)
}
}
impl TextBackend for FakeChatBackend {
fn tokenizer(&self) -> DynTokenizer {
Arc::new(FakeChatTokenizer)
Arc::new(TestTokenizer::new())
}
fn model_id(&self) -> &str {
@@ -194,11 +194,12 @@ mod tests {
use axum::http::HeaderMap;
use serde_json::json;
use vllm_text::Prompt;
use vllm_text::tokenizer::Tokenizer;
use vllm_tokenizer::test_utils::TestTokenizer;
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 {
@@ -212,32 +213,8 @@ mod tests {
}
}
#[derive(Debug)]
struct TestTokenizer;
impl Tokenizer for TestTokenizer {
fn encode(
&self,
text: &str,
_add_special_tokens: bool,
) -> vllm_text::tokenizer::Result<Vec<u32>> {
Ok(text.bytes().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_text::tokenizer::Result<String> {
Ok(
String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
.into_owned(),
)
}
fn token_to_id(&self, _token: &str) -> Option<u32> {
None
}
fn test_tokenizer() -> TestTokenizer {
TestTokenizer::new()
}
fn base_request_json() -> serde_json::Value {
@@ -273,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!({
@@ -297,7 +296,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare");
@@ -340,7 +339,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare")
.text_request
@@ -374,7 +373,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare");
@@ -399,7 +398,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare");
@@ -422,7 +421,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare");
@@ -446,7 +445,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare");
@@ -471,7 +470,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare");
@@ -498,7 +497,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare");
@@ -523,7 +522,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare");
@@ -549,7 +548,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare");
assert_eq!(prepared.text_request.sampling_params.logprobs, Some(1));
@@ -574,7 +573,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
request_context(&headers, None),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare");
assert_eq!(prepared.text_request.data_parallel_rank, Some(3));
@@ -593,7 +592,7 @@ mod tests {
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
&TestTokenizer,
&test_tokenizer(),
)
.expect("prepare");
assert_eq!(prepared.text_request.data_parallel_rank, None);
@@ -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]
+22 -145
View File
@@ -40,8 +40,9 @@ use vllm_engine_core_client::{
};
use vllm_llm::Llm;
use vllm_metrics::METRICS;
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
use vllm_text::tokenizer::DynTokenizer;
use vllm_text::{Prompt, TextBackend};
use vllm_tokenizer::test_utils::TestTokenizer;
use zeromq::prelude::{SocketRecv, SocketSend};
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
@@ -417,82 +418,20 @@ struct FakeChatBackend {
}
/// Synthetic BOS id used when `add_special_tokens` is true in tests.
const FAKE_BOS_TOKEN_ID: u32 = 1;
const FAKE_BOS_TOKEN_ID: u32 = 256;
const UNKNOWN_DECODE_TOKEN_ID: u32 = 10_000;
#[derive(Debug)]
struct FakeChatTokenizer;
impl Tokenizer for FakeChatTokenizer {
fn encode(
&self,
text: &str,
add_special_tokens: bool,
) -> vllm_text::tokenizer::Result<Vec<u32>> {
let mut token_ids = Vec::new();
if add_special_tokens {
token_ids.push(FAKE_BOS_TOKEN_ID);
}
let mut rest = text;
while !rest.is_empty() {
if let Some(stripped) = rest.strip_prefix("<image>") {
token_ids.push(999);
rest = stripped;
continue;
}
if let Some(stripped) = rest.strip_prefix("<|image_pad|>") {
token_ids.push(151655);
rest = stripped;
continue;
}
let ch = rest.chars().next().expect("rest is not empty");
let mut buf = [0; 4];
token_ids.extend(ch.encode_utf8(&mut buf).bytes().map(u32::from));
rest = &rest[ch.len_utf8()..];
}
Ok(token_ids)
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_text::tokenizer::Result<String> {
Ok(
String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
.into_owned(),
)
}
fn token_to_id(&self, token: &str) -> Option<u32> {
match token {
"<image>" => Some(999),
"<|image_pad|>" => Some(151655),
"<think>" => Some(0xF001),
"</think>" => Some(0xF002),
"<|START_THINKING|>" => Some(0xF003),
"<|END_THINKING|>" => Some(0xF004),
"◁think▷" => Some(0xF005),
"◁/think▷" => Some(0xF006),
_ => None,
}
}
fn id_to_token(&self, id: u32) -> Option<String> {
match id {
FAKE_BOS_TOKEN_ID => Some("<bos>".to_string()),
999 => Some("<image>".to_string()),
151655 => Some("<|image_pad|>".to_string()),
0xF001 => Some("<think>".to_string()),
0xF002 => Some("</think>".to_string()),
0xF003 => Some("<|START_THINKING|>".to_string()),
0xF004 => Some("<|END_THINKING|>".to_string()),
0xF005 => Some("◁think▷".to_string()),
0xF006 => Some("◁/think▷".to_string()),
id if id < 128 => char::from_u32(id).map(|ch| ch.to_string()),
_ => None,
}
}
fn fake_chat_tokenizer() -> TestTokenizer {
TestTokenizer::new()
.with_bos_token("<bos>", FAKE_BOS_TOKEN_ID)
.with_regular_token("<image>", 999)
.with_regular_token("<|image_pad|>", 151655)
.with_regular_token("<think>", 0xF001)
.with_regular_token("</think>", 0xF002)
.with_regular_token("<|START_THINKING|>", 0xF003)
.with_regular_token("<|END_THINKING|>", 0xF004)
.with_regular_token("◁think▷", 0xF005)
.with_regular_token("◁/think▷", 0xF006)
}
impl FakeChatBackend {
@@ -530,7 +469,7 @@ impl fmt::Debug for FakeChatBackend {
impl TextBackend for FakeChatBackend {
fn tokenizer(&self) -> DynTokenizer {
Arc::new(FakeChatTokenizer)
Arc::new(fake_chat_tokenizer())
}
fn model_id(&self) -> &str {
@@ -630,7 +569,7 @@ fn qwen_multimodal_model_info() -> vllm_chat::multimodal::MultimodalModelInfo {
Some("qwen2_vl".to_string()),
Some(&config_path),
None,
Arc::new(FakeChatTokenizer),
Arc::new(fake_chat_tokenizer()),
)
.expect("load multimodal info")
.expect("qwen multimodal info is registered");
@@ -638,70 +577,6 @@ fn qwen_multimodal_model_info() -> vllm_chat::multimodal::MultimodalModelInfo {
info
}
#[derive(Clone, Debug)]
struct FailingDecodeChatBackend;
#[derive(Debug)]
struct FailingDecodeTokenizer;
impl Tokenizer for FailingDecodeTokenizer {
fn encode(
&self,
text: &str,
add_special_tokens: bool,
) -> vllm_text::tokenizer::Result<Vec<u32>> {
FakeChatTokenizer.encode(text, add_special_tokens)
}
fn decode(
&self,
token_ids: &[u32],
skip_special_tokens: bool,
) -> vllm_text::tokenizer::Result<String> {
if token_ids.contains(&(b'i' as u32)) {
return Err(vllm_text::tokenizer::TokenizerError(
"forced decode failure for streaming test".to_string(),
));
}
FakeChatTokenizer.decode(token_ids, skip_special_tokens)
}
fn token_to_id(&self, token: &str) -> Option<u32> {
FakeChatTokenizer.token_to_id(token)
}
}
impl TextBackend for FailingDecodeChatBackend {
fn tokenizer(&self) -> DynTokenizer {
Arc::new(FailingDecodeTokenizer)
}
fn model_id(&self) -> &str {
"test-model"
}
}
impl ChatBackend for FailingDecodeChatBackend {
fn chat_renderer(&self) -> DynChatRenderer {
Arc::new(self.clone())
}
fn new_chat_output_processor(
&self,
_request: &mut ChatRequest,
_options: NewChatOutputProcessorOptions<'_>,
) -> vllm_chat::Result<DynChatOutputProcessor> {
Ok(Box::new(DefaultChatOutputProcessor::plain_text_only()))
}
}
impl ChatRenderer for FailingDecodeChatBackend {
fn render(&self, request: &ChatRequest) -> vllm_chat::Result<vllm_chat::RenderedPrompt> {
FakeChatBackend::new().render(request)
}
}
async fn test_models_with_engine_outputs_and_backend_inner(
engine_id: impl Into<EngineId>,
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
@@ -2755,8 +2630,8 @@ async fn load_endpoint_resets_when_stream_response_is_dropped() {
#[serial]
async fn stream_error_is_returned_as_openai_error_sse() {
let (app, engine_task) = test_app_with_backend_and_stream_output_specs(
Arc::new(FailingDecodeChatBackend),
default_stream_output_specs(),
Arc::new(FakeChatBackend::new()),
vec![(vec![UNKNOWN_DECODE_TOKEN_ID], None)],
)
.await;
let response = app
@@ -2789,7 +2664,9 @@ async fn stream_error_is_returned_as_openai_error_sse() {
assert!(text.contains("\"role\":\"assistant\""), "{text}");
assert!(text.contains("\"type\":\"server_error\""), "{text}");
assert!(
text.contains("forced decode failure for streaming test"),
text.contains(&format!(
"test tokenizer cannot decode unknown token id {UNKNOWN_DECODE_TOKEN_ID}"
)),
"{text}"
);
assert!(!text.contains("\"usage\":"), "{text}");
+18 -29
View File
@@ -1,5 +1,5 @@
//! TLS tests: `build_server_config` unit checks plus end-to-end OpenSSL handshakes
//! through the production `serve_listener` path, with a trivial router since TLS
//! through the production listener/connection path, with a trivial router since TLS
//! terminates below the app.
use std::pin::Pin;
@@ -23,8 +23,8 @@ use tokio_openssl::SslStream;
use tokio_util::sync::CancellationToken;
use crate::config::{HttpListenerMode, TlsConfig};
use crate::listener::Listener;
use crate::{ConnectionTimeouts, serve_listener, tls};
use crate::listener::{Listener, MaybeTlsListener};
use crate::{ConnectionTimeouts, serve_connections, tls};
// ============================================================================
// Test infrastructure
@@ -251,7 +251,6 @@ fn build_tls(certs: &TestCerts, cert: &str, key: Option<&str>) -> TlsConfig {
/// Generous per-connection timeouts that never fire during the fast tests.
const TEST_TIMEOUTS: ConnectionTimeouts = ConnectionTimeouts {
handshake: Duration::from_secs(60),
header_read: Duration::from_secs(5),
keep_alive_enabled: true,
};
@@ -261,9 +260,8 @@ async fn spawn_server(tls_config: Option<TlsConfig>) -> (String, CancellationTok
}
/// Bind an ephemeral listener and serve a trivial router via the production
/// `serve_listener`, optionally with TLS. The listener is bound (and thus
/// accepting into the backlog) before returning, so a client may connect
/// immediately without a sleep.
/// listener/connection path. The listener is bound (and thus accepting into the
/// backlog) before returning, so a client may connect immediately without a sleep.
async fn spawn_server_with_timeouts(
tls_config: Option<TlsConfig>,
timeouts: ConnectionTimeouts,
@@ -274,7 +272,7 @@ async fn spawn_server_with_timeouts(
})
.await
.expect("bind listener");
let addr = listener.local_addr().expect("local addr");
let addr = listener.local_addr_display().expect("local addr");
let server_config =
tls_config.map(|cfg| tls::build_server_config(&cfg).expect("build server config"));
@@ -282,14 +280,11 @@ async fn spawn_server_with_timeouts(
let shutdown = CancellationToken::new();
let server_shutdown = shutdown.clone();
tokio::spawn(async move {
let _ = serve_listener(
listener,
server_config,
app,
server_shutdown.cancelled_owned(),
timeouts,
)
.await;
let listener = match server_config {
Some(context) => MaybeTlsListener::tls(listener, context),
None => MaybeTlsListener::plain(listener),
};
let _ = serve_connections(listener, app, server_shutdown.cancelled_owned(), timeouts).await;
});
(addr, shutdown)
}
@@ -521,20 +516,19 @@ async fn plain_http_serves_when_tls_is_disabled() {
shutdown.cancel();
}
#[tokio::test]
#[tokio::test(start_paused = true)]
async fn tls_handshake_timeout_drops_silent_client() {
// Silent client (no ClientHello) must be dropped at the handshake deadline.
let certs = TestCerts::generate();
let timeouts = ConnectionTimeouts {
handshake: Duration::from_millis(150),
header_read: Duration::from_secs(5),
keep_alive_enabled: true,
};
let (addr, shutdown) = spawn_server_with_timeouts(Some(server_tls(&certs, 0)), timeouts).await;
let (addr, shutdown) = spawn_server(Some(server_tls(&certs, 0))).await;
let mut tcp = TcpStream::connect(&addr).await.expect("connect");
tokio::task::yield_now().await;
tokio::time::advance(tls::TLS_HANDSHAKE_TIMEOUT + Duration::from_millis(1)).await;
tokio::task::yield_now().await;
let mut buf = [0u8; 1];
let read = tokio::time::timeout(Duration::from_secs(5), tcp.read(&mut buf)).await;
let read = tokio::time::timeout(Duration::from_secs(1), tcp.read(&mut buf)).await;
assert!(
matches!(read, Ok(Ok(0)) | Ok(Err(_))),
"server must drop a stalled TLS handshake (expected close, got {read:?})"
@@ -546,7 +540,6 @@ async fn tls_handshake_timeout_drops_silent_client() {
async fn keep_alive_timeout_closes_idle_connection() {
// Idle keep-alive connection must be closed at the deadline.
let timeouts = ConnectionTimeouts {
handshake: Duration::from_secs(60),
header_read: Duration::from_millis(150),
keep_alive_enabled: true,
};
@@ -582,7 +575,6 @@ async fn keep_alive_timeout_closes_idle_tls_connection() {
// still fires through tls-listener's post-handshake SslStream, not just plaintext.
let certs = TestCerts::generate();
let timeouts = ConnectionTimeouts {
handshake: Duration::from_secs(60),
header_read: Duration::from_millis(150),
keep_alive_enabled: true,
};
@@ -618,7 +610,6 @@ async fn keep_alive_timeout_closes_idle_tls_connection() {
async fn idle_timeout_closes_silent_client() {
// Silent client closed by the header-read timeout (http1-only arms it from byte 0).
let timeouts = ConnectionTimeouts {
handshake: Duration::from_secs(60),
header_read: Duration::from_millis(150),
keep_alive_enabled: true,
};
@@ -638,7 +629,6 @@ async fn idle_timeout_closes_silent_client() {
async fn keep_alive_zero_disables_keep_alive() {
// 0 disables keep-alive (serve, then close), like uvicorn's timeout_keep_alive=0.
let timeouts = ConnectionTimeouts {
handshake: Duration::from_secs(60),
header_read: Duration::from_secs(5),
keep_alive_enabled: false,
};
@@ -671,7 +661,6 @@ async fn disabled_keep_alive_still_closes_silent_client() {
// Even with keep-alive off, the head read stays bounded, so a silent client
// is dropped rather than held open.
let timeouts = ConnectionTimeouts {
handshake: Duration::from_secs(60),
header_read: Duration::from_millis(150),
keep_alive_enabled: false,
};
+2 -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]
@@ -31,6 +31,7 @@ serial_test.workspace = true
tempfile.workspace = true
tokio.workspace = true
vllm-llm = { workspace = true, features = ["test-util"] }
vllm-tokenizer = { workspace = true, features = ["test-utils"] }
[lints]
workspace = true
+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}"
+4 -57
View File
@@ -275,6 +275,7 @@ mod tests {
use std::collections::{BTreeSet, HashMap};
use serial_test::file_serial;
use vllm_tokenizer::test_utils::TestTokenizer;
use super::*;
use crate::backend::hf::HfTextBackend;
@@ -282,60 +283,8 @@ mod tests {
use crate::error::{LogprobsError, TokenIdsError};
use crate::request::{Prompt, TextRequest};
/// Stub tokenizer that returns empty token IDs — sufficient for tests that
/// don't exercise bad-words tokenization.
struct StubTokenizer;
impl Tokenizer for StubTokenizer {
fn encode(
&self,
_text: &str,
_add_special_tokens: bool,
) -> vllm_tokenizer::Result<Vec<u32>> {
Ok(vec![])
}
fn decode(
&self,
_token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok(String::new())
}
fn token_to_id(&self, _token: &str) -> Option<u32> {
None
}
}
fn stub_tokenizer() -> StubTokenizer {
StubTokenizer
}
struct FixedTokenizer {
token_ids: Vec<u32>,
}
impl Tokenizer for FixedTokenizer {
fn encode(
&self,
_text: &str,
_add_special_tokens: bool,
) -> vllm_tokenizer::Result<Vec<u32>> {
Ok(self.token_ids.clone())
}
fn decode(
&self,
_token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok(String::new())
}
fn token_to_id(&self, _token: &str) -> Option<u32> {
None
}
fn stub_tokenizer() -> TestTokenizer {
TestTokenizer::new()
}
fn sample_request() -> TextRequest {
@@ -952,9 +901,7 @@ mod tests {
#[test]
fn lower_sampling_params_rejects_out_of_vocab_bad_words() {
let tokenizer = FixedTokenizer {
token_ids: vec![1999, 2000],
};
let tokenizer = TestTokenizer::new().with_regular_token("blocked", 2000);
let error = lower_sampling_params(
SamplingParams {
bad_words: Some(vec!["blocked".to_string()]),
+3 -29
View File
@@ -323,37 +323,11 @@ mod tests {
use futures::{Stream, stream};
use vllm_engine_core_client::AbortCause;
use vllm_llm::GenerateOutput;
use vllm_tokenizer::Tokenizer;
use vllm_tokenizer::test_utils::TestTokenizer;
use super::*;
use crate::output::TextOutputStreamExt as _;
/// Backend that treats each token ID as a raw byte, producing lossy UTF-8.
struct ByteTokenizer;
impl Tokenizer for ByteTokenizer {
fn encode(
&self,
_text: &str,
_add_special_tokens: bool,
) -> vllm_tokenizer::Result<Vec<u32>> {
unreachable!()
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
let bytes = token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>();
Ok(String::from_utf8_lossy(&bytes).into_owned())
}
fn token_to_id(&self, _token: &str) -> Option<u32> {
unreachable!()
}
}
/// Helper: run `decoded_text_event_stream` to completion and return the
/// collected output.
async fn run_to_completion(
@@ -366,7 +340,7 @@ mod tests {
token_ids,
Some(FinishReason::Length),
))]);
let tokenizer: DynTokenizer = Arc::new(ByteTokenizer);
let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
decoded_text_event_stream("test".into(), tokenizer, raw_stream, decode_options, false)
.collect_output()
.await
@@ -419,7 +393,7 @@ mod tests {
))),
dropped_cause: Arc::clone(&dropped_cause),
};
let tokenizer: DynTokenizer = Arc::new(ByteTokenizer);
let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
let output = decoded_text_event_stream(
"test".into(),
+3 -30
View File
@@ -129,40 +129,13 @@ fn decode_position_logprobs<T: Tokenizer + ?Sized>(
#[cfg(test)]
mod tests {
use vllm_llm::{Logprobs, PositionLogprobs, TokenLogprob};
use vllm_tokenizer::test_utils::TestTokenizer;
use super::*;
#[derive(Debug)]
struct ByteTokenizer;
impl vllm_tokenizer::Tokenizer for ByteTokenizer {
fn encode(
&self,
_text: &str,
_add_special_tokens: bool,
) -> vllm_tokenizer::Result<Vec<u32>> {
unreachable!()
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok(String::from_utf8_lossy(
&token_ids.iter().map(|token_id| *token_id as u8).collect::<Vec<_>>(),
)
.into_owned())
}
fn token_to_id(&self, _token: &str) -> Option<u32> {
unreachable!()
}
}
#[test]
fn decode_logprobs_decodes_every_candidate_token() {
let tokenizer = ByteTokenizer;
let tokenizer = TestTokenizer::new();
let logprobs = Logprobs {
positions: vec![PositionLogprobs {
entries: vec![
@@ -205,7 +178,7 @@ mod tests {
#[test]
fn decode_prompt_logprobs_separates_first_prompt_token() {
let tokenizer = ByteTokenizer;
let tokenizer = TestTokenizer::new();
let logprobs = Logprobs {
positions: vec![PositionLogprobs {
entries: vec![TokenLogprob {
+3
View File
@@ -4,6 +4,9 @@ version.workspace = true
edition.workspace = true
license.workspace = true
[features]
test-utils = []
[dependencies]
base64.workspace = true
fastokens.workspace = true
+12
View File
@@ -170,6 +170,10 @@ mod tests {
fn token_to_id(&self, _token: &str) -> Option<u32> {
unreachable!()
}
fn id_to_token(&self, _id: u32) -> Option<String> {
unreachable!()
}
}
#[test]
@@ -248,6 +252,10 @@ mod tests {
fn token_to_id(&self, _token: &str) -> Option<u32> {
unreachable!()
}
fn id_to_token(&self, _id: u32) -> Option<String> {
unreachable!()
}
}
#[test]
@@ -320,6 +328,10 @@ mod tests {
fn token_to_id(&self, _token: &str) -> Option<u32> {
unreachable!()
}
fn id_to_token(&self, _id: u32) -> Option<String> {
unreachable!()
}
}
/// Without the char-boundary fix, this panics slicing mid-emoji.
+3 -5
View File
@@ -8,6 +8,8 @@ mod error;
mod hf;
mod incremental;
mod tekken;
#[cfg(any(test, feature = "test-utils"))]
pub mod test_utils;
mod tiktoken;
pub use error::{Result, TokenizerError};
@@ -28,11 +30,7 @@ pub trait Tokenizer: Send + Sync {
fn token_to_id(&self, token: &str) -> Option<u32>;
/// Convert one token ID into the tokenizer's raw token string.
fn id_to_token(&self, _id: u32) -> Option<String> {
// TODO: remove default impl and require this to be implemented by all
// tokenizers
None
}
fn id_to_token(&self, id: u32) -> Option<String>;
/// Return the vocabulary size. Backends that cannot report it fall back to
/// `usize::MAX`, an effectively unbounded value used only by test stubs.

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