Compare commits

...
Author SHA1 Message Date
Bugen Zhao cf30ef60cf use newtype for logprobs count (-1 for all)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-22 19:30:10 +08:00
Tuukka SarviandGitHub 89accad2cc [ROCm][DSV4] Disable TileLang MHC dispatch on gfx942 (#45931)
Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com>
2026-06-22 09:26:54 +00:00
3c8e49596c [Model] ColQwen3.5: fix retrieval correctness (bias + bidirectional) (#46108)
Signed-off-by: Athrael Soju <athrael.soju@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-22 17:25:54 +08:00
Weiwei SunGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Jiangyun Zhu
cec2ec1176 [Bugfix] Avoid racy accepted counts in async spec decode (#45100)
Signed-off-by: Weiwei Sun <68775773+sunnweiwei@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
2026-06-22 08:53:16 +00:00
liuzhenweiandGitHub 435f82d61a [Bugfix] Fix Llama4ForCausalLM initialization test failure (#46341)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
2026-06-22 08:40:43 +00:00
Roger WangandGitHub 1c4b51b990 Temporarily skip M3 on CI (#46352)
Signed-off-by: Roger Wang <hey@rogerw.io>
2026-06-22 01:35:31 -07:00
2e2c47928b [Doc] Update MiniMax-M3 (#45940)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Signed-off-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
2026-06-22 01:23:27 -07:00
Chao-Ju ChenGitHubBugen Zhaomergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
80abe0de7d [Rust Frontend] Support thinking_token_budget for chat and completions (#46137)
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: RickyChen / 陳昭儒 <ricky.chen@infinirc.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-22 16:00:02 +08:00
a9f7b2d41c [feature][kv_offload] Self-describing KV events for OffloadingConnector (#43468)
Signed-off-by: Change72 <changg@nvidia.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-06-22 07:27:46 +00:00
Tiezhen WANGGitHubClaudemergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
d14e551a53 [Model] Remove MiniMaxText01, MiniMaxVL01, MiniMaxForCausalLM (#45993)
Signed-off-by: Xianbao QIAN <xianbao.qian@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-22 15:20:46 +08:00
68567ef2df [CPUOffloadingManager] Maintain evictable list in LRUCachePolicy (#46216)
Signed-off-by: <>
Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>
2026-06-22 06:54:44 +00:00
6bc6f2d86d [1/N][Core] add partial prefix cache primitives (#45939)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-06-21 23:43:10 -07:00
wang.yuqiandGitHub 1eb2cc961e [Frontend] Refactor ServingTokenization entrypoint. (#46022)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-06-22 06:27:58 +00:00
ReidGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
31124749d1 [Bugfix] [Rust Frontend] Fix stop string truncation with repeated matches (#46113)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: reidliu41 <reid201711@gmail.com>
2026-06-22 14:11:29 +08:00
Ma JianandGitHub 9037498c22 [DSV4][XPU] Pass gemm1_clamp_limit to XpuFusedMoe (#44517)
Signed-off-by: Ma Jian <jian1.ma@intel.com>
2026-06-22 12:57:10 +08:00
db32b53e30 [SpecDecode] Support DFlash with FlashInfer (#43081)
Signed-off-by: gss <2783977641@qq.com>
Co-authored-by: gss <2783977641@qq.com>
2026-06-22 04:55:30 +00:00
xiangdongandGitHub b529bfd6c5 [XPU][CI] Add agent_tags for Intel GPU CI (#45768)
Signed-off-by: zengxian <xiangdong.zeng@intel.com>
2026-06-22 10:33:17 +08:00
Micah WilliamsonandGitHub f3df7a7231 [ROCm][CI] Enable kv_connector unit tests on ROCm (#45955)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-06-22 05:08:44 +03:00
485bbe1c6f [CI] Fix missing tp_size attribute on RoutedExperts (#46163)
Signed-off-by: Felix Marty <Felix.Marty@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-06-21 18:46:49 -06:00
114 changed files with 2707 additions and 4529 deletions
@@ -21,6 +21,10 @@ steps:
timeout_in_minutes: 30
optional: true
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 24+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -38,6 +42,10 @@ steps:
timeout_in_minutes: 30
optional: true
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -55,6 +63,10 @@ steps:
timeout_in_minutes: 30
optional: true
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -5,6 +5,10 @@ steps:
- label: XPU Sleep Mode
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
+4
View File
@@ -5,6 +5,10 @@ steps:
- label: Engine (1 GPU)
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -6,6 +6,10 @@ steps:
key: eplb-algorithm
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
+4
View File
@@ -5,6 +5,10 @@ steps:
- label: vLLM IR Tests
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
+24
View File
@@ -5,6 +5,10 @@ steps:
- label: LoRA Runtime + Utils
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
@@ -34,6 +38,10 @@ steps:
- label: LoRA Fused/MoE Kernels
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -54,6 +62,10 @@ steps:
- label: LoRA Punica Kernels
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -74,6 +86,10 @@ steps:
- label: LoRA Punica FP8/XPU Ops
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -94,6 +110,10 @@ steps:
- label: LoRA Models
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 24+
no_plugin: true
working_dir: "."
env:
@@ -117,6 +137,10 @@ steps:
- label: LoRA Multimodal
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
+24
View File
@@ -5,6 +5,10 @@ steps:
- label: V1 Core + KV + Metrics
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -31,6 +35,10 @@ steps:
- label: V1 Sample + Logits
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -71,6 +79,10 @@ steps:
- label: XPU CPU Offload
timeout_in_minutes: 60
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -95,6 +107,10 @@ steps:
key: regression
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -126,6 +142,10 @@ steps:
timeout_in_minutes: 30
num_devices: 2
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -157,6 +177,10 @@ steps:
key: async-engine-inputs-utils-worker
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
@@ -5,6 +5,10 @@ steps:
- label: Model Runner V2 Core Tests (Intel)
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -30,6 +34,10 @@ steps:
- label: Model Runner V2 Examples (Intel)
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
@@ -6,6 +6,10 @@ steps:
key: multi-modal-models-standard-1-qwen2
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -27,6 +31,10 @@ steps:
key: multi-modal-models-standard-2-qwen3-gemma
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -47,6 +55,10 @@ steps:
key: multi-modal-models-standard-3-llava-qwen2-vl
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
@@ -68,6 +80,10 @@ steps:
key: multi-modal-models-standard-4-other-whisper
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
@@ -88,6 +104,10 @@ steps:
key: multi-modal-processor
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
+16
View File
@@ -19,6 +19,10 @@ steps:
- image-build-xpu
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 24+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -49,6 +53,10 @@ steps:
- image-build-xpu
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -74,6 +82,10 @@ steps:
- image-build-xpu
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -93,6 +105,10 @@ steps:
- image-build-xpu
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -4,6 +4,11 @@
set -euo pipefail
if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None else 1)"; then
uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
exit 0
fi
REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
uv pip install --system -r "${REQUIREMENTS_FILE}"
+6
View File
@@ -105,6 +105,12 @@ steps:
# Integration test for streaming correctness (requires special branch).
- pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
mirror:
amd:
device: mi325_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
- label: V1 Others (CPU)
key: v1-others-cpu
+2 -2
View File
@@ -136,7 +136,7 @@ The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/mo
For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together).
These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol).
For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively.
For case (3), we recommend looking at the implementation of [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which uses a custom "mamba-like" layer `ShortConv`.
Please follow the same guidelines as case (2) for implementing these models.
We use "mamba-like" to refer to layers that possess a state that is updated in-place, rather than being appended-to (like KV cache for attention).
For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`.
@@ -144,5 +144,5 @@ It is also necessary to implement the "attention meta-data" class which handles
Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.
It is also worth noting that we should update `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/attention/backends/registry.py) when adding a new mamba backend.
Finally, if one wants to support torch compile and CUDA graphs, it necessary to wrap the call to the mamba-like layer inside a custom op and register it.
Please see the calls to `direct_register_custom_op` in [vllm/model_executor/models/minimax_text_01.py](../../../vllm/model_executor/models/minimax_text_01.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
Please see the calls to `direct_register_custom_op` in [vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py](../../../vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
The new custom op should then be added to the list `_attention_ops` in [vllm/config/compilation.py](../../../vllm/config/compilation.py) to ensure that piecewise CUDA graphs works as intended.
+2 -2
View File
@@ -170,8 +170,8 @@ Priority is **1 = highest** (tried first).
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
+1
View File
@@ -74,6 +74,7 @@ vllm serve <model> \
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
## Secondary Tiers
-9
View File
@@ -321,15 +321,6 @@ For Qwen2.5, the chat template in tokenizer_config.json has already included sup
Flags: `--tool-call-parser hermes`
### MiniMax Models (`minimax_m1`)
Supported models:
* `MiniMaxAi/MiniMax-M1-40k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
* `MiniMaxAi/MiniMax-M1-80k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
Flags: `--tool-call-parser minimax --chat-template examples/tool_chat_template_minimax_m1.jinja`
### DeepSeek-V3 Models (`deepseek_v3`)
Supported models:
+1 -1
View File
@@ -61,7 +61,7 @@ Models of any architecture can be converted into embedding models using `--conve
| `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
| `ColPaliForRetrieval` | ColPali | T / I | `vidore/colpali-v1.3-hf` | | |
| `ColQwen3` | Qwen3-VL | T / I | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` | | |
| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3` | | |
| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3`, `vultr/VultronRetrieverPrime-Qwen3.5-8B` | | |
| `OpsColQwen3Model` | Qwen3-VL | T / I | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` | | |
| `Qwen3VLNemotronEmbedModel` | Qwen3-VL | T / I | `nvidia/nemotron-colembed-vl-4b-v2`, `nvidia/nemotron-colembed-vl-8b-v2` | ✅︎ | ✅︎ |
| `*ForConditionalGeneration`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | \* | N/A | \* | \* |
+1 -3
View File
@@ -441,7 +441,6 @@ th {
| `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
| `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 | `MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ |
| `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ |
| `MistralLarge3ForCausalLM` | Mistral-Large-3-675B-Base-2512, Mistral-Large-3-675B-Instruct-2512 | `mistralai/Mistral-Large-3-675B-Base-2512`, `mistralai/Mistral-Large-3-675B-Instruct-2512`, etc. | ✅︎ | ✅︎ |
@@ -487,8 +486,6 @@ th {
| `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
| `TeleChat3ForCausalLM` | TeleChat3 | `Tele-AI/TeleChat3-36B-Thinking`, `Tele-AI/TeleChat3-Coder-36B-Thinking`, etc. | ✅︎ | ✅︎ |
| `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
| `MiniMaxM1ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`, etc. | | |
| `MiniMaxText01ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01`, etc. | | |
| `Zamba2ForCausalLM` | Zamba2 | `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc. | | |
!!! note
@@ -595,6 +592,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
| `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
| `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I<sup>+</sup> + V<sup>+</sup> | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | |
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
| `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
| `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
+1 -1
View File
@@ -128,7 +128,7 @@ Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaFor
Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`,
`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`).
Hybrid models with mechanisms different to Mamba are also supported (e.g, `MiniMaxText01ForCausalLM`, `MiniMaxM1ForCausalLM`, `Lfm2ForCausalLM`).
Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`).
Please note that prefix caching is not yet supported for any of the above models.
@@ -1481,39 +1481,6 @@ def run_minicpmv(questions: list[str], modality: str) -> ModelRequestData:
return run_minicpmv_base(questions, modality, "openbmb/MiniCPM-V-2_6")
def run_minimax_vl_01(questions: list[str], modality: str) -> ModelRequestData:
assert modality == "image"
model_name = "MiniMaxAI/MiniMax-VL-01"
engine_args = EngineArgs(
model=model_name,
max_num_seqs=2,
limit_mm_per_prompt={modality: 1},
trust_remote_code=True,
tensor_parallel_size=8,
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
messages = [
[
{
"role": "user",
"content": [{"type": "image"}, {"type": "text", "text": question}],
}
]
for question in questions
]
prompts = tokenizer.apply_chat_template(
messages, add_generation_prompt=True, tokenize=False
)
return ModelRequestData(
engine_args=engine_args,
prompts=prompts,
)
# Mistral-3 HF-format
def run_mistral3(questions: list[str], modality: str) -> ModelRequestData:
assert modality == "image"
@@ -2485,7 +2452,6 @@ model_example_map = {
"mantis": run_mantis,
"minicpmo": run_minicpmo,
"minicpmv": run_minicpmv,
"minimax_vl_01": run_minimax_vl_01,
"mistral3": run_mistral3,
"molmo": run_molmo,
"molmo2": run_molmo2,
@@ -7,11 +7,27 @@ ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5.
It produces per-token embeddings and uses MaxSim scoring for retrieval
and reranking. Supports both text and image inputs.
Works for any ColQwen3.5 checkpoint, e.g. `athrael-soju/colqwen3.5-4.5B-v3`
or `vultr/VultronRetrieverPrime-Qwen3.5-8B`.
Start the server with:
vllm serve athrael-soju/colqwen3.5-4.5B --max-model-len 4096
vllm serve athrael-soju/colqwen3.5-4.5B-v3 --max-model-len 4096 \
--mm-processor-kwargs '{"min_pixels": 65536, "max_pixels": 1835008}'
Then run this script:
python colqwen3_5_rerank_online.py
Parity note (matching the native colpali ColQwen3_5Processor pipeline):
- Visual-token budget: ColQwen3_5Processor uses max_num_visual_tokens=1792,
i.e. max_pixels = 1792 * (patch_size*merge_size)^2 = 1792 * 32^2 = 1835008
(with min_pixels = shortest_edge = 65536). Pass these via --mm-processor-kwargs
as above; the default budget gives fewer visual tokens and lower retrieval ndcg.
- When you build prompts yourself (token_embed), reproduce the processor exactly:
image (document): wrap in the instruction template
"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
"Describe the image.<|im_end|><|endoftext|>"
query: append the augmentation suffix <text> + "<|endoftext|>" * 10
Omitting these reproduces a silent ~2.5 ndcg@10 drop vs the native pipeline.
"""
import requests
@@ -1,91 +0,0 @@
{{ '<begin_of_document>' -}}
{%- if custom_tools is defined %}
{%- set tools = custom_tools %}
{%- endif %}
{%- if not tools is defined %}
{%- set tools = none %}
{%- endif %}
{#- Extract system message #}
{% set ns = namespace(system_prompt='') -%}
{%- if messages[0]['role'] == 'system' %}
{%- if messages[0]['content'] is string %}
{%- set ns.system_prompt = messages[0]['content']|trim %}
{%- else %}
{%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
{%- endif %}
{%- set messages = messages[1:] %}
{%- else %}
{%- if tools is not none %}
{%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
{%- else %}
{%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
{%- endif %}
{%- endif %}
{#- System message #}
{%- if ns.system_prompt != '' %}
{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
{%- endif %}
{#- Tools configuration #}
{%- if tools is not none %}
{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
{%- for tool in tools %}
{{ tool | tojson ~ '\n' -}}
{%- endfor %}
{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
{%- endif %}
{#- Process messages #}
{%- for message in messages %}
{%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
{%- if message['role'] == 'user' %}
{{ '<beginning_of_sentence>user name=user\n' -}}
{%- if message['content'] is string %}
{{ message['content']|trim -}}
{%- else %}
{%- for content in message['content'] %}
{%- if content['type'] == 'text' %}
{{ content['text']|trim -}}
{%- endif %}
{%- endfor %}
{%- endif %}
{{ '<end_of_sentence>\n' -}}
{%- elif message['role'] == 'assistant' %}
{{ '<beginning_of_sentence>ai name=assistant\n' -}}
{%- if message['content'] is string %}
{{ message['content']|trim -}}
{%- else %}
{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
{{ content['text']|trim -}}
{%- endfor %}
{%- endif %}
{{ '<end_of_sentence>\n' -}}
{%- endif %}
{%- elif 'tool_calls' in message %}
{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
{%- for tool_call in message.tool_calls %}
{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
{%- endfor %}
{{ '</tool_calls><end_of_sentence>\n' -}}
{%- elif message.role == "tool" or message.role == "ipython" %}
{{ '<beginning_of_sentence>tool name=tools\n' -}}
{%- if message.content is string %}
{{ 'tool result: ' + message.content + '\n\n' -}}
{%- else %}
{%- for content in message['content'] %}
{%- if content['type'] == 'text' %}
{{ 'tool result: ' + content['text'] + '\n\n' -}}
{%- elif content.get('name') %}
{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
{%- endif %}
{%- endfor %}
{%- endif %}
{{ '<end_of_sentence>\n' -}}
{%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
{{ '<beginning_of_sentence>ai name=assistant\n' -}}
{%- endif %}
-1
View File
@@ -386,7 +386,6 @@ mod tests {
tool_chat_template_llama3.2_pythonic.jinja => String
tool_chat_template_llama4_json.jinja => OpenAi
tool_chat_template_llama4_pythonic.jinja => OpenAi
tool_chat_template_minimax_m1.jinja => OpenAi
tool_chat_template_mistral.jinja => String
tool_chat_template_mistral3.jinja => OpenAi
tool_chat_template_mistral_parallel.jinja => String
+4 -3
View File
@@ -16,7 +16,8 @@ use vllm_engine_core_client::protocol::logprobs::{
Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
};
use vllm_engine_core_client::protocol::{
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, LogprobsCount,
StopReason,
};
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
@@ -1387,8 +1388,8 @@ async fn chat_stream_and_collect_preserve_prompt_and_sample_logprobs() {
.await;
let mut request = sample_request("chat-logprobs");
request.sampling_params.logprobs = Some(1);
request.sampling_params.prompt_logprobs = Some(1);
request.sampling_params.logprobs = Some(LogprobsCount::Top(1));
request.sampling_params.prompt_logprobs = Some(LogprobsCount::Top(1));
let mut stream = chat.chat(request.clone()).await.unwrap();
match next_semantic(&mut stream).await.unwrap().unwrap() {
@@ -1,91 +0,0 @@
{{ '<begin_of_document>' -}}
{%- if custom_tools is defined %}
{%- set tools = custom_tools %}
{%- endif %}
{%- if not tools is defined %}
{%- set tools = none %}
{%- endif %}
{#- Extract system message #}
{% set ns = namespace(system_prompt='') -%}
{%- if messages[0]['role'] == 'system' %}
{%- if messages[0]['content'] is string %}
{%- set ns.system_prompt = messages[0]['content']|trim %}
{%- else %}
{%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
{%- endif %}
{%- set messages = messages[1:] %}
{%- else %}
{%- if tools is not none %}
{%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
{%- else %}
{%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
{%- endif %}
{%- endif %}
{#- System message #}
{%- if ns.system_prompt != '' %}
{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
{%- endif %}
{#- Tools configuration #}
{%- if tools is not none %}
{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
{%- for tool in tools %}
{{ tool | tojson ~ '\n' -}}
{%- endfor %}
{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
{%- endif %}
{#- Process messages #}
{%- for message in messages %}
{%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
{%- if message['role'] == 'user' %}
{{ '<beginning_of_sentence>user name=user\n' -}}
{%- if message['content'] is string %}
{{ message['content']|trim -}}
{%- else %}
{%- for content in message['content'] %}
{%- if content['type'] == 'text' %}
{{ content['text']|trim -}}
{%- endif %}
{%- endfor %}
{%- endif %}
{{ '<end_of_sentence>\n' -}}
{%- elif message['role'] == 'assistant' %}
{{ '<beginning_of_sentence>ai name=assistant\n' -}}
{%- if message['content'] is string %}
{{ message['content']|trim -}}
{%- else %}
{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
{{ content['text']|trim -}}
{%- endfor %}
{%- endif %}
{{ '<end_of_sentence>\n' -}}
{%- endif %}
{%- elif 'tool_calls' in message %}
{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
{%- for tool_call in message.tool_calls %}
{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
{%- endfor %}
{{ '</tool_calls><end_of_sentence>\n' -}}
{%- elif message.role == "tool" or message.role == "ipython" %}
{{ '<beginning_of_sentence>tool name=tools\n' -}}
{%- if message.content is string %}
{{ 'tool result: ' + message.content + '\n\n' -}}
{%- else %}
{%- for content in message['content'] %}
{%- if content['type'] == 'text' %}
{{ 'tool result: ' + content['text'] + '\n\n' -}}
{%- elif content.get('name') %}
{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
{%- endif %}
{%- endfor %}
{%- endif %}
{{ '<end_of_sentence>\n' -}}
{%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
{{ '<beginning_of_sentence>ai name=assistant\n' -}}
{%- endif %}
+8 -3
View File
@@ -20,6 +20,7 @@ use serde_with::{DefaultOnNull, OneOrMany, serde_as};
use thiserror_ext::AsReport as _;
use uuid::Uuid;
use vllm_engine_core_client::TransportMode;
use vllm_engine_core_client::protocol::LogprobsCount;
use vllm_managed_engine::ManagedEngineConfig;
use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args};
use vllm_server::{
@@ -136,9 +137,9 @@ pub struct SharedRuntimeArgs {
pub max_model_len: Option<u32>,
/// Maximum number of log probabilities to return when `logprobs` is
/// specified in sampling parameters. `-1` means no cap.
#[arg(long, value_parser = clap::value_parser!(i32).range(-1..), allow_negative_numbers = true)]
#[arg(long, allow_negative_numbers = true)]
#[serde(default)]
pub max_logprobs: Option<i32>,
pub max_logprobs: Option<LogprobsCount>,
/// TCP port for the gRPC Generate service. When not set, no gRPC server is
/// started.
#[arg(long)]
@@ -529,7 +530,7 @@ impl ServeArgs {
self.managed_engine.clone().into_config(
self.runtime.model.clone(),
self.runtime.max_model_len,
self.runtime.max_logprobs,
self.runtime.max_logprobs.map(managed_max_logprobs_to_i32),
self.runtime.language_model_only,
self.runtime.disable_log_stats,
self.runtime.shutdown_timeout,
@@ -555,5 +556,9 @@ fn frontend_ipc_addresses() -> (String, String) {
)
}
fn managed_max_logprobs_to_i32(count: LogprobsCount) -> i32 {
i32::try_from(count).expect("max_logprobs is parsed through i32")
}
#[cfg(test)]
mod tests;
+4 -3
View File
@@ -1,5 +1,6 @@
use expect_test::expect;
use vllm_engine_core_client::TransportMode;
use vllm_engine_core_client::protocol::LogprobsCount;
use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
use super::{Cli, Command};
@@ -165,10 +166,10 @@ fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() {
let Command::Serve(args) = cli.command else {
panic!("expected serve args");
};
assert_eq!(args.runtime.max_logprobs, Some(-1));
assert_eq!(args.runtime.max_logprobs, Some(LogprobsCount::All));
let frontend_config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
assert_eq!(frontend_config.max_logprobs, Some(-1));
assert_eq!(frontend_config.max_logprobs, Some(LogprobsCount::All));
let engine_config = args.to_managed_engine_config(5555);
assert_eq!(engine_config.python_args, vec!["--max-logprobs", "-1"]);
@@ -529,7 +530,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
assert!(args.runtime.language_model_only);
assert_eq!(args.runtime.max_model_len, Some(8192));
assert_eq!(args.runtime.max_logprobs, Some(-1));
assert_eq!(args.runtime.max_logprobs, Some(LogprobsCount::All));
assert_eq!(args.runtime.shutdown_timeout, 3);
}
@@ -6,7 +6,7 @@ use futures::StreamExt as _;
use tokio::time::timeout;
use tracing_subscriber::EnvFilter;
use vllm_engine_core_client::protocol::{
EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams, LogprobsCount,
};
use vllm_engine_core_client::{
EngineCoreClient, EngineCoreClientConfig, EngineCoreStreamOutput, TransportMode,
@@ -33,10 +33,10 @@ struct Args {
output_timeout_secs: u64,
#[arg(long, default_value_t = 1)]
max_tokens: u32,
#[arg(long, default_value_t = 2)]
logprobs: i32,
#[arg(long, default_value_t = 1)]
prompt_logprobs: i32,
#[arg(long, default_value_t = LogprobsCount::Top(2), allow_negative_numbers = true)]
logprobs: LogprobsCount,
#[arg(long, default_value_t = LogprobsCount::Top(1), allow_negative_numbers = true)]
prompt_logprobs: LogprobsCount,
#[arg(long, default_value_t = 96)]
prompt_repeats: usize,
}
@@ -64,8 +64,8 @@ fn build_request(
request_id: String,
prompt_token_ids: Vec<u32>,
max_tokens: u32,
logprobs: i32,
prompt_logprobs: i32,
logprobs: LogprobsCount,
prompt_logprobs: LogprobsCount,
client_index: u32,
) -> EngineCoreRequest {
EngineCoreRequest {
@@ -0,0 +1,135 @@
use std::fmt;
use std::str::FromStr;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
/// Number of log probabilities requested for a token position.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LogprobsCount {
/// Return the full model vocabulary.
All,
/// Return the top-N tokens by probability.
Top(u32),
}
impl LogprobsCount {
/// Expands the count to the actual number of logprobs to return, given the vocabulary size.
pub fn expanded(self, vocab_size: usize) -> usize {
match self {
Self::All => vocab_size,
Self::Top(count) => count as usize,
}
}
}
impl TryFrom<i32> for LogprobsCount {
type Error = String;
fn try_from(value: i32) -> Result<Self, Self::Error> {
match value {
-1 => Ok(Self::All),
value if value < -1 => Err(format!("must be non-negative or -1, got {value}")),
value => Ok(Self::Top(value as u32)),
}
}
}
impl TryFrom<LogprobsCount> for i32 {
type Error = String;
fn try_from(value: LogprobsCount) -> Result<Self, Self::Error> {
match value {
LogprobsCount::All => Ok(-1),
LogprobsCount::Top(count) => {
i32::try_from(count).map_err(|_| format!("must fit within i32, got {count}"))
}
}
}
}
impl FromStr for LogprobsCount {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let value = s
.parse::<i32>()
.map_err(|e| format!("must be an i32 integer, got {s:?}: {e}"))?;
Self::try_from(value)
}
}
impl fmt::Display for LogprobsCount {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::All => (-1).fmt(f),
Self::Top(count) => count.fmt(f),
}
}
}
impl Serialize for LogprobsCount {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let value: i32 = (*self).try_into().map_err(serde::ser::Error::custom)?;
value.serialize(serializer)
}
}
impl<'de> Deserialize<'de> for LogprobsCount {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let value = i32::deserialize(deserializer)?;
Self::try_from(value).map_err(serde::de::Error::custom)
}
}
#[cfg(test)]
mod tests {
use rmpv::Value;
use super::*;
use crate::protocol::{decode_msgpack, encode_msgpack};
#[test]
fn logprobs_count_serializes_as_wire_integer() {
assert_eq!(serde_json::to_value(LogprobsCount::All).unwrap(), -1);
assert_eq!(serde_json::to_value(LogprobsCount::Top(3)).unwrap(), 3);
}
#[test]
fn logprobs_count_deserializes_wire_integer() {
assert_eq!(
serde_json::from_value::<LogprobsCount>(serde_json::json!(-1)).unwrap(),
LogprobsCount::All
);
assert_eq!(
serde_json::from_value::<LogprobsCount>(serde_json::json!(3)).unwrap(),
LogprobsCount::Top(3)
);
assert!(serde_json::from_value::<LogprobsCount>(serde_json::json!(-2)).is_err());
assert!(
serde_json::from_value::<LogprobsCount>(serde_json::json!(i64::from(i32::MAX) + 1))
.is_err()
);
}
#[test]
fn logprobs_count_decodes_msgpack_signed_and_unsigned() {
let mut encoded = Vec::new();
rmpv::encode::write_value(&mut encoded, &Value::from(-1)).unwrap();
assert_eq!(
decode_msgpack::<LogprobsCount>(&encoded).unwrap(),
LogprobsCount::All
);
let encoded = encode_msgpack(&LogprobsCount::Top(7)).unwrap();
assert_eq!(
decode_msgpack::<LogprobsCount>(&encoded).unwrap(),
LogprobsCount::Top(7)
);
}
}
@@ -56,6 +56,7 @@ mod classified_outputs;
pub mod dtype;
pub mod handshake;
pub mod logprobs;
mod logprobs_count;
pub mod lora;
pub mod multimodal;
pub mod stats;
@@ -66,6 +67,7 @@ pub use classified_outputs::{
};
pub use dtype::ModelDtype;
pub use logprobs::decode_engine_core_outputs;
pub use logprobs_count::LogprobsCount;
/// Request types are encoded as single-byte protocol constants so they can be
/// sent over the ZMQ socket without an extra encoding step.
@@ -277,14 +279,20 @@ pub struct EngineCoreSamplingParams {
pub max_tokens: u32,
/// Minimum number of tokens to generate before EOS or stop-token handling.
pub min_tokens: u32,
/// Maximum number of reasoning ("thinking") tokens to emit before the
/// reasoning section is force-closed. `None` means unlimited; the
/// user-facing `-1` sentinel is normalized to `None` by the frontend before
/// reaching this DTO, so only non-negative values are sent. Enforced
/// engine-side (and only when a reasoning parser is configured).
pub thinking_token_budget: Option<u64>,
/// Number of log probabilities to return per generated token.
///
/// `None` disables sample logprobs. `-1` requests the full vocabulary.
pub logprobs: Option<i32>,
/// `None` disables sample logprobs.
pub logprobs: Option<LogprobsCount>,
/// Number of log probabilities to return per prompt token.
///
/// `None` disables prompt logprobs. `-1` requests the full vocabulary.
pub prompt_logprobs: Option<i32>,
/// `None` disables prompt logprobs.
pub prompt_logprobs: Option<LogprobsCount>,
/// Minimum probability threshold for token sampling.
pub min_p: f32,
/// Frequency penalty applied by the sampler.
@@ -345,6 +353,7 @@ impl EngineCoreSamplingParams {
seed: None,
max_tokens: 65536,
min_tokens: 0,
thinking_token_budget: None,
logprobs: None,
prompt_logprobs: None,
min_p: 0.0,
@@ -150,6 +150,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {
top_k: 8,
max_tokens: 32,
min_tokens: 1,
thinking_token_budget: Some(256),
stop_token_ids: vec![151643],
eos_token_id: Some(151645),
all_stop_token_ids: BTreeSet::from([151643, 151645]),
@@ -2502,6 +2503,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
seed: None,
max_tokens: 16,
min_tokens: 0,
thinking_token_budget: None,
logprobs: None,
prompt_logprobs: None,
min_p: 0.0,
@@ -39,6 +39,7 @@ class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True):
seed: int | None = None
max_tokens: int = 16
min_tokens: int = 0
thinking_token_budget: int | None = None
min_p: float = 0.0
frequency_penalty: float = 0.0
presence_penalty: float = 0.0
@@ -122,6 +123,7 @@ request = EngineCoreRequest(
seed=None,
max_tokens=32,
min_tokens=1,
thinking_token_budget=256,
min_p=0.0,
frequency_penalty=0.0,
presence_penalty=0.0,
+3 -11
View File
@@ -2,12 +2,13 @@ use std::collections::HashMap;
use std::fmt;
use std::time::Duration;
use anyhow::{Result, bail};
use anyhow::Result;
use axum::http::{HeaderName, HeaderValue, Method};
use educe::Educe;
use serde::Serialize;
use serde_json::Value;
use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
use vllm_engine_core_client::protocol::LogprobsCount;
use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode};
/// How the HTTP server obtains its listening socket.
@@ -133,7 +134,7 @@ pub struct Config {
pub chat_template_content_format: ChatTemplateContentFormatOption,
/// Optional maximum number of top log probabilities accepted by the
/// frontend. `None` delegates to the text layer default.
pub max_logprobs: Option<i32>,
pub max_logprobs: Option<LogprobsCount>,
/// HTTP/API-server behavior switches.
pub api_server_options: ApiServerOptions,
/// CORS settings applied to every HTTP response.
@@ -158,15 +159,6 @@ impl Config {
pub fn validate(&self) -> Result<()> {
vllm_chat::validate_parser_overrides(&self.tool_call_parser, &self.reasoning_parser)?;
self.cors.validate()?;
if let Some(max_logprobs) = self.max_logprobs
&& max_logprobs < -1
{
bail!(
"max_logprobs must be non-negative or -1, got {}",
max_logprobs
);
}
Ok(())
}
+13
View File
@@ -103,6 +103,7 @@ fn is_request_validation_error(error: &vllm_text::Error) -> bool {
| vllm_text::Error::EmptyPromptTokenIds { .. }
| vllm_text::Error::Logprobs(_)
| vllm_text::Error::OutOfVocab(_)
| vllm_text::Error::InvalidThinkingTokenBudget
// An empty tokenized prompt detected later, at request prepare
// time, surfaces through the transparent Llm wrapper.
| vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds { .. })
@@ -127,6 +128,18 @@ mod tests {
assert!(response.error.message.contains("9000"));
}
#[test]
fn invalid_thinking_token_budget_maps_to_invalid_request() {
let api_error = text_submit_error(
"failed to submit completion request",
vllm_text::Error::InvalidThinkingTokenBudget,
);
assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
let response = api_error.to_error_response();
assert_eq!(response.error.error_type, "invalid_request_error");
assert!(response.error.message.contains("thinking_token_budget"));
}
#[test]
fn chat_wrapped_prompt_too_long_maps_to_invalid_request() {
let error = vllm_chat::Error::Text(vllm_text::Error::PromptTooLong {
+13 -9
View File
@@ -3,7 +3,7 @@
use tonic::Status;
use uuid::Uuid;
use vllm_engine_core_client::protocol::{StopReason, StructuredOutputsParams};
use vllm_engine_core_client::protocol::{LogprobsCount, StopReason, StructuredOutputsParams};
use vllm_text::{
DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, SamplingParams,
TextDecodeOptions, TextRequest,
@@ -202,18 +202,22 @@ fn build_sampling_params(
/// Map the proto `CandidateTokens` selector to a `(logprobs_count,
/// logprob_token_ids)` pair.
///
/// - `top_n(k)` → `(k, None)` — return top-k candidates by probability
/// - `all` → `(-1, None)` — return the full vocabulary
/// - `top_n(k)` → `(Top(k), None)` — return top-k candidates by probability
/// - `all` → `(All, None)` — return the full vocabulary
/// - `token_ids(n)` → `(1, Some(vec of n token ids))` — return logprobs for specific tokens (the
/// count `n` is stored in the proto as the number of token IDs that follow, but the actual IDs
/// are carried via `logprob_token_ids` on `SamplingParams`)
/// - absent → `(1, None)` — just the sampled/scored token
fn candidate_logprob_spec(candidates: Option<&pb::CandidateTokens>) -> (i32, Option<Vec<u32>>) {
/// - absent → `(Top(1), None)` — just the sampled/scored token
fn candidate_logprob_spec(
candidates: Option<&pb::CandidateTokens>,
) -> (LogprobsCount, Option<Vec<u32>>) {
match candidates.and_then(|c| c.select.as_ref()) {
Some(pb::candidate_tokens::Select::TopN(n)) => (*n as i32, None),
Some(pb::candidate_tokens::Select::All(true)) => (-1, None),
Some(pb::candidate_tokens::Select::TokenIds(ids)) => (1, Some(ids.ids.clone())),
_ => (1, None),
Some(pb::candidate_tokens::Select::TopN(n)) => (LogprobsCount::Top(*n), None),
Some(pb::candidate_tokens::Select::All(true)) => (LogprobsCount::All, None),
Some(pb::candidate_tokens::Select::TokenIds(ids)) => {
(LogprobsCount::Top(1), Some(ids.ids.clone()))
}
_ => (LogprobsCount::Top(1), None),
}
}
@@ -87,6 +87,7 @@ pub(super) fn prepare_generate_request(
#[cfg(test)]
mod tests {
use serde_json::json;
use vllm_engine_core_client::protocol::LogprobsCount;
use vllm_text::Prompt;
use super::prepare_generate_request;
@@ -132,10 +133,13 @@ mod tests {
Prompt::TokenIds(vec![11, 22, 33])
);
assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7));
assert_eq!(prepared.text_request.sampling_params.logprobs, Some(2));
assert_eq!(
prepared.text_request.sampling_params.logprobs,
Some(LogprobsCount::Top(2))
);
assert_eq!(
prepared.text_request.sampling_params.prompt_logprobs,
Some(1)
Some(LogprobsCount::Top(1))
);
assert!(prepared.text_request.sampling_params.ignore_eos);
assert_eq!(prepared.text_request.priority, -3);
@@ -150,6 +154,33 @@ mod tests {
);
}
#[test]
fn prepare_generate_request_forwards_thinking_token_budget() {
let request: GenerateRequest = serde_json::from_value(json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"token_ids": [11, 22, 33],
"sampling_params": {
"thinking_token_budget": 64
}
}))
.expect("parse request");
let prepared = prepare_generate_request(
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
)
.expect("prepare");
// The raw inference route shares `vllm_text::SamplingParams`, so the
// field is carried through to lowering exactly like the OpenAI routes
// (normalization/validation then happens in `lower_sampling_params`).
assert_eq!(
prepared.text_request.sampling_params.thinking_token_budget,
Some(64)
);
}
#[test]
fn prepare_generate_request_gates_continuous_usage_on_include_usage() {
let request: GenerateRequest = serde_json::from_value(json!({
@@ -34,16 +34,6 @@ pub(super) fn validate_request_compat(
);
}
if let Some(prompt_logprobs) = request.sampling_params.prompt_logprobs
&& prompt_logprobs < 0
&& prompt_logprobs != -1
{
bail_invalid_request!(
param = "sampling_params",
"`prompt_logprobs` must be a non-negative value or -1."
);
}
Ok(())
}
@@ -4,6 +4,7 @@ use vllm_chat::{
ChatMessage as VllmChatMessage, ChatOptions, ChatRequest, ChatTool, ChatToolChoice,
GenerationPromptMode, SamplingParams,
};
use vllm_engine_core_client::protocol::LogprobsCount;
use super::types::ChatCompletionRequest;
use super::validate;
@@ -94,7 +95,7 @@ pub(super) fn prepare_chat_request(
// Auto-enable prompt logprobs for non-streaming echo, matching Python vLLM's
// behavior.
let top_logprobs = request.top_logprobs.unwrap_or(0);
let top_logprobs = request.top_logprobs.unwrap_or(LogprobsCount::Top(0));
let prompt_logprobs = request
.prompt_logprobs
.or((request.echo && !request.stream).then_some(top_logprobs));
@@ -115,6 +116,7 @@ pub(super) fn prepare_chat_request(
seed: request.seed,
max_tokens: request.max_completion_tokens,
min_tokens: request.min_tokens,
thinking_token_budget: request.thinking_token_budget,
logprobs: request.logprobs.then_some(top_logprobs),
prompt_logprobs,
min_p: request.min_p,
@@ -377,6 +379,7 @@ mod tests {
ChatTool as VllmChatTool, ChatToolChoice, GenerationPromptMode,
SamplingParams as VllmSamplingParams,
};
use vllm_engine_core_client::protocol::LogprobsCount;
use vllm_text::output::TextDecodeOptions;
use super::prepare_chat_request;
@@ -613,6 +616,31 @@ mod tests {
assert_eq!(prepared.chat_request.sampling_params, expected);
}
#[test]
fn prepare_chat_request_passes_through_thinking_token_budget() {
let prepare = |budget: Option<i64>| {
prepare_chat_request(
ChatCompletionRequest {
thinking_token_budget: budget,
..base_request()
},
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
)
.expect("request is valid")
.chat_request
.sampling_params
.thinking_token_budget
};
// The convert layer forwards the raw value verbatim (including the `-1`
// "unlimited" sentinel); normalization/validation happens during
// lowering (see `vllm_text::lower`).
assert_eq!(prepare(Some(64)), Some(64));
assert_eq!(prepare(Some(-1)), Some(-1));
assert_eq!(prepare(None), None);
}
#[test]
fn prepare_chat_request_accepts_developer_messages() {
let request = ChatCompletionRequest {
@@ -941,7 +969,7 @@ mod tests {
let request = ChatCompletionRequest {
stream: false,
logprobs: true,
prompt_logprobs: Some(2),
prompt_logprobs: Some(LogprobsCount::Top(2)),
..base_request()
};
@@ -954,10 +982,13 @@ mod tests {
assert!(prepared.options.requested_logprobs);
assert!(prepared.options.include_prompt_logprobs);
assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(0));
assert_eq!(
prepared.chat_request.sampling_params.logprobs,
Some(LogprobsCount::Top(0))
);
assert_eq!(
prepared.chat_request.sampling_params.prompt_logprobs,
Some(2)
Some(LogprobsCount::Top(2))
);
}
@@ -965,7 +996,7 @@ mod tests {
fn prepare_chat_request_keeps_prompt_logprobs_independent_from_echo() {
let request = ChatCompletionRequest {
logprobs: true,
top_logprobs: Some(3),
top_logprobs: Some(LogprobsCount::Top(3)),
echo: true,
..base_request()
};
@@ -977,7 +1008,10 @@ mod tests {
)
.expect("request is valid");
assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(3));
assert_eq!(
prepared.chat_request.sampling_params.logprobs,
Some(LogprobsCount::Top(3))
);
assert_eq!(prepared.chat_request.sampling_params.prompt_logprobs, None);
assert!(!prepared.options.include_prompt_logprobs);
}
@@ -6,6 +6,7 @@ use serde_json::Value;
use serde_with::SerializeDisplay;
use validator::Validate;
use vllm_chat::ReasoningEffort;
use vllm_engine_core_client::protocol::LogprobsCount;
use crate::routes::openai::utils::structured_outputs::ResponseFormat;
use crate::routes::openai::utils::types::{
@@ -44,10 +45,8 @@ pub struct ChatCompletionRequest {
#[serde(default)]
pub logprobs: bool,
/// An integer specifying the number of most likely tokens to return
/// -1 means return all
#[validate(range(min = -1))]
pub top_logprobs: Option<i32>,
/// Number of most likely tokens to return. `-1` means return full vocab.
pub top_logprobs: Option<LogprobsCount>,
/// Deprecated: Replaced by max_completion_tokens
#[deprecated(note = "Use max_completion_tokens instead")]
@@ -155,8 +154,8 @@ pub struct ChatCompletionRequest {
/// Truncate prompt tokens to this length
pub truncate_prompt_tokens: Option<i64>,
/// Number of prompt logprobs to return
pub prompt_logprobs: Option<i32>,
/// Number of prompt logprobs to return. `-1` means return full vocab.
pub prompt_logprobs: Option<LogprobsCount>,
/// Restrict output to these token IDs only
pub allowed_token_ids: Option<Vec<u32>>,
@@ -165,8 +164,10 @@ pub struct ChatCompletionRequest {
pub bad_words: Option<Vec<String>>,
// -------- Extra vLLM Parameters --------
/// Token budget for reasoning/thinking
pub thinking_token_budget: Option<u32>,
/// Token budget for reasoning/thinking. Accepts a non-negative integer, or
/// `-1` for unlimited (mirroring the Python frontend, which normalizes `-1`
/// to "no budget").
pub thinking_token_budget: Option<i64>,
/// Whether to include reasoning content in the response
#[serde(default = "default_true")]
@@ -1,6 +1,7 @@
use super::types::ChatCompletionRequest;
use crate::error::{ApiError, bail_invalid_request};
use crate::routes::openai::utils::types::{ChatMessage, Tool, ToolChoice, ToolChoiceValue};
use vllm_engine_core_client::protocol::LogprobsCount;
/// Enforce the minimal compatibility contract for the Rust OpenAI server.
pub(super) fn validate_request_compat(
@@ -30,14 +31,12 @@ pub(super) fn validate_request_compat(
}
if let Some(prompt_logprobs) = request.prompt_logprobs {
if prompt_logprobs < 0 && prompt_logprobs != -1 {
bail_invalid_request!(
param = "prompt_logprobs",
"prompt_logprobs must be a non-negative value or -1."
);
}
if request.stream && (prompt_logprobs > 0 || prompt_logprobs == -1) {
if request.stream
&& matches!(
prompt_logprobs,
LogprobsCount::All | LogprobsCount::Top(1..)
)
{
bail_invalid_request!(
param = "prompt_logprobs",
"prompt_logprobs are not available when stream=true."
@@ -108,11 +107,6 @@ pub(super) fn validate_request_compat(
"truncate_prompt_tokens",
"truncate_prompt_tokens is not supported.",
)?;
reject_non_default(
request.thinking_token_budget.as_ref(),
"thinking_token_budget",
"thinking_token_budget is not supported.",
)?;
reject_non_default(
request.media_io_kwargs.as_ref(),
"media_io_kwargs",
@@ -159,6 +153,7 @@ mod tests {
use serde_json::json;
use vllm_chat::ReasoningEffort;
use vllm_engine_core_client::protocol::LogprobsCount;
use super::validate_request_compat;
use crate::routes::openai::chat_completions::types::ChatCompletionRequest;
@@ -304,7 +299,7 @@ mod tests {
#[test]
fn validate_request_compat_rejects_top_logprobs_without_logprobs() {
let request = ChatCompletionRequest {
top_logprobs: Some(0),
top_logprobs: Some(LogprobsCount::Top(0)),
..base_request()
};
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
@@ -313,26 +308,26 @@ mod tests {
#[test]
fn validate_request_compat_rejects_streaming_prompt_logprobs_requests() {
let request = ChatCompletionRequest {
prompt_logprobs: Some(1),
prompt_logprobs: Some(LogprobsCount::Top(1)),
..base_request()
};
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
let request = ChatCompletionRequest {
prompt_logprobs: Some(-1),
prompt_logprobs: Some(LogprobsCount::All),
..base_request()
};
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
}
#[test]
fn validate_request_compat_rejects_invalid_prompt_logprobs_value() {
let request = ChatCompletionRequest {
stream: false,
prompt_logprobs: Some(-2),
..base_request()
};
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
fn chat_request_deserialization_rejects_invalid_prompt_logprobs_value() {
let result = serde_json::from_value::<ChatCompletionRequest>(json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"messages": [{"role": "user", "content": "hello"}],
"prompt_logprobs": -2
}));
assert!(result.is_err());
}
#[test]
@@ -1,3 +1,4 @@
use vllm_engine_core_client::protocol::LogprobsCount;
use vllm_text::{SamplingParams, TextDecodeOptions, TextRequest};
use super::types::CompletionRequest;
@@ -61,15 +62,7 @@ pub(super) fn prepare_completion_request(
.map(|request| request.lora_name.clone())
.unwrap_or_else(|| lora_resolution.model_names.first().cloned().unwrap_or_default());
let logprobs = match request.logprobs {
Some(logprobs) => Some(i32::try_from(logprobs).map_err(|_| {
ApiError::invalid_request(
"`logprobs` must fit within a signed 32-bit integer.".to_string(),
Some("logprobs"),
)
})?),
None => None,
};
let logprobs = request.logprobs.map(LogprobsCount::Top);
let prompt_only = request.echo && request.max_tokens == Some(0);
let prompt_logprobs =
request.prompt_logprobs.or(if request.echo && (!request.stream || prompt_only) {
@@ -108,6 +101,7 @@ pub(super) fn prepare_completion_request(
seed: request.seed,
max_tokens,
min_tokens: request.min_tokens,
thinking_token_budget: request.thinking_token_budget,
logprobs,
prompt_logprobs,
min_p: request.min_p,
@@ -162,6 +156,7 @@ pub(super) fn prepare_completion_request(
mod tests {
use axum::http::HeaderMap;
use serde_json::json;
use vllm_engine_core_client::protocol::LogprobsCount;
use vllm_text::Prompt;
use super::prepare_completion_request;
@@ -246,7 +241,10 @@ mod tests {
Prompt::TokenIds(vec![11, 22, 33])
);
assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7));
assert_eq!(prepared.text_request.sampling_params.logprobs, Some(2));
assert_eq!(
prepared.text_request.sampling_params.logprobs,
Some(LogprobsCount::Top(2))
);
assert_eq!(prepared.text_request.sampling_params.top_p, Some(0.9));
assert_eq!(prepared.text_request.sampling_params.top_k, Some(42));
assert_eq!(prepared.text_request.sampling_params.min_p, Some(0.1));
@@ -266,6 +264,34 @@ mod tests {
assert!(!prepared.text_request.decode_options.skip_special_tokens);
}
#[test]
fn prepare_completion_request_passes_through_thinking_token_budget() {
let prepare = |budget: serde_json::Value| {
let request: CompletionRequest = serde_json::from_value(json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"prompt": "hello",
"thinking_token_budget": budget,
}))
.expect("parse request");
prepare_completion_request(
request,
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
ResolvedRequestContext::default(),
)
.expect("prepare")
.text_request
.sampling_params
.thinking_token_budget
};
// The convert layer forwards the raw value verbatim (including the `-1`
// "unlimited" sentinel); normalization/validation happens during
// lowering (see `vllm_text::lower`).
assert_eq!(prepare(json!(64)), Some(64));
assert_eq!(prepare(json!(-1)), Some(-1));
assert_eq!(prepare(json!(null)), None);
}
#[test]
fn prepare_completion_request_maps_stream_usage_and_token_format_options() {
let request: CompletionRequest = serde_json::from_value(json!({
@@ -381,10 +407,13 @@ mod tests {
.expect("prepare");
assert!(prepared.options.prompt_only);
assert_eq!(prepared.text_request.sampling_params.logprobs, Some(3));
assert_eq!(
prepared.text_request.sampling_params.logprobs,
Some(LogprobsCount::Top(3))
);
assert_eq!(
prepared.text_request.sampling_params.prompt_logprobs,
Some(3)
Some(LogprobsCount::Top(3))
);
}
@@ -406,10 +435,13 @@ mod tests {
)
.expect("prepare");
assert_eq!(prepared.text_request.sampling_params.logprobs, Some(3));
assert_eq!(
prepared.text_request.sampling_params.logprobs,
Some(LogprobsCount::Top(3))
);
assert_eq!(
prepared.text_request.sampling_params.prompt_logprobs,
Some(3)
Some(LogprobsCount::Top(3))
);
}
@@ -450,10 +482,13 @@ mod tests {
ResolvedRequestContext::default(),
)
.expect("prepare");
assert_eq!(prepared.text_request.sampling_params.logprobs, Some(1));
assert_eq!(
prepared.text_request.sampling_params.logprobs,
Some(LogprobsCount::Top(1))
);
assert_eq!(
prepared.text_request.sampling_params.prompt_logprobs,
Some(2)
Some(LogprobsCount::Top(2))
);
}
@@ -3,6 +3,7 @@ use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use validator::Validate;
use vllm_engine_core_client::protocol::LogprobsCount;
use vllm_text::Prompt;
use crate::routes::openai::utils::types::{
@@ -131,8 +132,8 @@ pub struct CompletionRequest {
/// Restrict output to these token IDs only
pub allowed_token_ids: Option<Vec<u32>>,
/// Number of prompt logprobs to return
pub prompt_logprobs: Option<i32>,
/// Number of prompt logprobs to return. `-1` means return full vocab.
pub prompt_logprobs: Option<LogprobsCount>,
// -------- Extra vLLM Parameters --------
/// Whether to add special tokens (e.g. BOS) to the prompt
@@ -146,6 +147,11 @@ pub struct CompletionRequest {
/// Additional kwargs for structured outputs
pub structured_outputs: Option<Value>,
/// Token budget for reasoning/thinking. Accepts a non-negative integer, or
/// `-1` for unlimited (mirroring the Python frontend, which normalizes `-1`
/// to "no budget").
pub thinking_token_budget: Option<i64>,
/// Request scheduling priority (lower means earlier; default 0)
pub priority: Option<i32>,
@@ -1,3 +1,4 @@
use vllm_engine_core_client::protocol::LogprobsCount;
use vllm_text::Prompt;
use super::types::CompletionRequest;
@@ -44,29 +45,18 @@ pub(super) fn validate_request_compat(
bail_invalid_request!(param = "suffix", "suffix is not supported.");
}
if let Some(logprobs) = request.logprobs
&& logprobs > i32::MAX as u32
{
bail_invalid_request!(
param = "logprobs",
"`logprobs` must fit within a signed 32-bit integer."
);
}
if let Some(prompt_logprobs) = request.prompt_logprobs {
if request.stream && (prompt_logprobs > 0 || prompt_logprobs == -1) {
if request.stream
&& matches!(
prompt_logprobs,
LogprobsCount::All | LogprobsCount::Top(1..)
)
{
bail_invalid_request!(
param = "prompt_logprobs",
"`prompt_logprobs` are not available when `stream=true`."
);
}
if prompt_logprobs < 0 && prompt_logprobs != -1 {
bail_invalid_request!(
param = "prompt_logprobs",
"`prompt_logprobs` must be a non-negative value or -1."
);
}
}
if request.use_beam_search {
@@ -101,6 +91,7 @@ pub(super) fn validate_request_compat(
#[cfg(test)]
mod tests {
use serde_json::json;
use vllm_engine_core_client::protocol::LogprobsCount;
use super::validate_request_compat;
use crate::routes::openai::completions::types::CompletionRequest;
@@ -150,7 +141,7 @@ mod tests {
#[test]
fn validate_request_compat_rejects_streaming_prompt_logprobs() {
let request = CompletionRequest {
prompt_logprobs: Some(1),
prompt_logprobs: Some(LogprobsCount::Top(1)),
..base_request()
};
assert!(
@@ -162,7 +153,7 @@ mod tests {
fn validate_request_compat_accepts_non_stream_prompt_logprobs() {
let request = CompletionRequest {
stream: false,
prompt_logprobs: Some(-1),
prompt_logprobs: Some(LogprobsCount::All),
..base_request()
};
assert!(
+3 -4
View File
@@ -2,6 +2,7 @@ pub mod hf;
use std::sync::Arc;
use vllm_engine_core_client::protocol::LogprobsCount;
use vllm_tokenizer::DynTokenizer;
use crate::error::Result;
@@ -26,9 +27,7 @@ pub struct SamplingLimits {
/// Runtime context window size reported by the engine startup handshake.
pub max_model_len: u32,
/// Maximum number of top log probabilities accepted by this frontend.
///
/// `-1` means allowing requests up to the model vocabulary size.
pub max_logprobs: i32,
pub max_logprobs: LogprobsCount,
/// Model vocabulary size from the model config, used to bound generated
/// token IDs and logits-domain sampling controls.
@@ -41,7 +40,7 @@ pub struct SamplingLimits {
impl SamplingLimits {
/// Original Python definition:
/// <https://github.com/vllm-project/vllm/blob/b5adb027ad03c29b46181752ba3b1cb84eff1dd4/vllm/config/model.py#L216-L220>
pub const DEFAULT_MAX_LOGPROBS: i32 = 20;
pub const DEFAULT_MAX_LOGPROBS: LogprobsCount = LogprobsCount::Top(20);
/// Original Python definition:
/// <https://github.com/vllm-project/vllm/blob/b5adb027ad03c29b46181752ba3b1cb84eff1dd4/vllm/sampling_params.py#L30-L32>
pub const MAX_LOGPROB_TOKEN_IDS: usize = 128;
+2
View File
@@ -20,6 +20,8 @@ pub enum Error {
Logprobs(#[from] LogprobsError),
#[error(transparent)]
OutOfVocab(#[from] OutOfVocabError),
#[error("`thinking_token_budget` must be a non-negative integer or -1 for unlimited.")]
InvalidThinkingTokenBudget,
#[error("text request stream `{request_id}` closed before terminal output")]
StreamClosedBeforeTerminalOutput { request_id: String },
#[error(transparent)]
+3 -2
View File
@@ -19,6 +19,7 @@ pub use output::{
pub use request::{Prompt, SamplingParams, TextRequest};
use trait_set::trait_set;
use vllm_engine_core_client::EngineCoreClient;
use vllm_engine_core_client::protocol::LogprobsCount;
pub use vllm_llm::FinishReason;
use vllm_llm::{GenerateOutputStream, Llm};
use vllm_tokenizer::DynTokenizer;
@@ -48,7 +49,7 @@ pub struct TextLlm {
/// Runtime context window size reported by the engine startup handshake.
max_model_len: u32,
/// Maximum number of top log probabilities accepted by this text facade.
max_logprobs: i32,
max_logprobs: LogprobsCount,
}
impl TextLlm {
@@ -68,7 +69,7 @@ impl TextLlm {
}
/// Override the maximum accepted logprobs count.
pub fn with_max_logprobs(mut self, max_logprobs: Option<i32>) -> Self {
pub fn with_max_logprobs(mut self, max_logprobs: Option<LogprobsCount>) -> Self {
if let Some(max_logprobs) = max_logprobs {
self.max_logprobs = max_logprobs;
}
+65 -10
View File
@@ -87,6 +87,7 @@ pub fn lower_sampling_params(
seed,
max_tokens,
min_tokens,
thinking_token_budget,
logprobs,
prompt_logprobs,
min_p,
@@ -128,6 +129,7 @@ pub fn lower_sampling_params(
prompt_len,
)?;
let min_tokens = min_tokens.unwrap_or(0);
let thinking_token_budget = normalize_thinking_token_budget(thinking_token_budget)?;
let frequency_penalty = frequency_penalty.unwrap_or(0.0);
let presence_penalty = presence_penalty.unwrap_or(0.0);
@@ -149,6 +151,7 @@ pub fn lower_sampling_params(
seed,
max_tokens,
min_tokens,
thinking_token_budget,
logprobs,
prompt_logprobs,
min_p,
@@ -170,6 +173,21 @@ pub fn lower_sampling_params(
Ok(params)
}
/// Normalize the user-facing `thinking_token_budget` into the engine value.
///
/// Mirrors Python's `validate_thinking_token_budget`
/// (<https://github.com/vllm-project/vllm/blob/ecf9d83520eb217401b47d8a5451a27c5231b8c2/vllm/sampling_params.py#L35-L55>):
/// `None` and the `-1` "unlimited" sentinel both map to `None`; any other
/// negative value is rejected; non-negative values pass through unchanged. Like
/// Python's `int`, no upper bound is imposed.
fn normalize_thinking_token_budget(value: Option<i64>) -> Result<Option<u64>> {
match value {
None | Some(-1) => Ok(None),
Some(budget) if budget >= 0 => Ok(Some(budget as u64)),
Some(_) => Err(Error::InvalidThinkingTokenBudget),
}
}
/// Convert bad-word strings into token-ID sequences, following the Python vLLM
/// logic in `SamplingParams.update_from_tokenizer()`.
///
@@ -251,6 +269,7 @@ mod tests {
use std::collections::{BTreeSet, HashMap};
use serial_test::file_serial;
use vllm_engine_core_client::protocol::LogprobsCount;
use super::*;
use crate::backend::hf::HfTextBackend;
@@ -366,6 +385,36 @@ mod tests {
)
}
#[test]
fn lower_sampling_params_normalizes_thinking_token_budget() {
let lower = |budget: Option<i64>| {
lower_sampling_params_with_limits(
SamplingParams {
thinking_token_budget: budget,
..SamplingParams::default()
},
sample_sampling_limits(),
)
};
// Non-negative budgets (including 0) pass through unchanged.
assert_eq!(lower(Some(256)).unwrap().thinking_token_budget, Some(256));
assert_eq!(lower(Some(0)).unwrap().thinking_token_budget, Some(0));
// `None` and the `-1` "unlimited" sentinel both disable the budget.
assert_eq!(lower(None).unwrap().thinking_token_budget, None);
assert_eq!(lower(Some(-1)).unwrap().thinking_token_budget, None);
// No upper bound is imposed, matching Python's `int`.
assert_eq!(
lower(Some(i64::from(u32::MAX) + 1)).unwrap().thinking_token_budget,
Some(u64::from(u32::MAX) + 1)
);
// Other negatives are rejected.
assert!(matches!(
lower(Some(-2)),
Err(Error::InvalidThinkingTokenBudget)
));
}
#[test]
fn lower_text_request_applies_python_style_eos_hints() {
let prepared = lower_text_request(
@@ -386,6 +435,7 @@ mod tests {
seed: None,
max_tokens: 999997,
min_tokens: 0,
thinking_token_budget: None,
logprobs: None,
prompt_logprobs: None,
min_p: 0.0,
@@ -437,6 +487,7 @@ mod tests {
seed: None,
max_tokens: 999997,
min_tokens: 0,
thinking_token_budget: None,
logprobs: None,
prompt_logprobs: None,
min_p: 0.0,
@@ -567,6 +618,7 @@ mod tests {
seed: None,
max_tokens: 40957,
min_tokens: 0,
thinking_token_budget: None,
logprobs: None,
prompt_logprobs: None,
min_p: 0.0,
@@ -628,6 +680,7 @@ mod tests {
seed: None,
max_tokens: 999997,
min_tokens: 0,
thinking_token_budget: None,
logprobs: None,
prompt_logprobs: None,
min_p: 0.0,
@@ -697,6 +750,7 @@ mod tests {
seed: None,
max_tokens: 32,
min_tokens: 2,
thinking_token_budget: None,
logprobs: None,
prompt_logprobs: None,
min_p: 0.1,
@@ -721,8 +775,8 @@ mod tests {
#[test]
fn lower_sampling_params_passes_logprobs_fields_through() {
let sampling_params = SamplingParams {
logprobs: Some(3),
prompt_logprobs: Some(-1),
logprobs: Some(LogprobsCount::Top(3)),
prompt_logprobs: Some(LogprobsCount::All),
..Default::default()
};
@@ -739,7 +793,7 @@ mod tests {
default_max_tokens: None,
},
SamplingLimits {
max_logprobs: -1,
max_logprobs: LogprobsCount::All,
..sample_sampling_limits()
},
3,
@@ -747,15 +801,15 @@ mod tests {
)
.unwrap();
assert_eq!(params.logprobs, Some(3));
assert_eq!(params.prompt_logprobs, Some(-1));
assert_eq!(params.logprobs, Some(LogprobsCount::Top(3)));
assert_eq!(params.prompt_logprobs, Some(LogprobsCount::All));
}
#[test]
fn lower_sampling_params_rejects_full_vocab_logprobs_over_default_cap() {
let error = lower_sampling_params_with_limits(
SamplingParams {
logprobs: Some(-1),
logprobs: Some(LogprobsCount::All),
..Default::default()
},
sample_sampling_limits(),
@@ -776,24 +830,24 @@ mod tests {
fn lower_sampling_params_expands_full_vocab_logprobs_from_model_vocab() {
let params = lower_sampling_params_with_limits(
SamplingParams {
logprobs: Some(-1),
logprobs: Some(LogprobsCount::All),
..Default::default()
},
SamplingLimits {
max_logprobs: 1500,
max_logprobs: LogprobsCount::Top(1500),
..sample_sampling_limits()
},
)
.unwrap();
assert_eq!(params.logprobs, Some(-1));
assert_eq!(params.logprobs, Some(LogprobsCount::All));
}
#[test]
fn lower_sampling_params_rejects_invalid_logprob_token_ids() {
let error = lower_sampling_params_with_limits(
SamplingParams {
logprobs: Some(1),
logprobs: Some(LogprobsCount::Top(1)),
logprob_token_ids: Some(vec![1000]),
..Default::default()
},
@@ -929,6 +983,7 @@ mod tests {
seed: None,
max_tokens: 128,
min_tokens: 0,
thinking_token_budget: None,
logprobs: None,
prompt_logprobs: None,
min_p: 0.1,
+13 -24
View File
@@ -1,15 +1,14 @@
//! Python-compatible validation for logprobs sampling params.
//!
//! `-1` is expanded only for bounds checks. The original request values are
//! `All` is expanded only for bounds checks. The original request values are
//! passed through to engine-core.
use crate::backend::SamplingLimits;
use thiserror::Error;
use vllm_engine_core_client::protocol::LogprobsCount;
#[derive(Debug, Error)]
pub enum LogprobsError {
#[error("{parameter} must be non-negative or -1, got {value}")]
InvalidCount { parameter: &'static str, value: i32 },
#[error(
"requested {parameter} of {requested}, which is greater than max allowed: {max_allowed}"
)]
@@ -30,19 +29,21 @@ pub enum LogprobsError {
"when both logprobs and logprob_token_ids are set, logprobs must equal \
len(logprob_token_ids). Got logprobs={logprobs}, len(logprob_token_ids)={num_token_ids}."
)]
TokenIdsMismatch { logprobs: i32, num_token_ids: usize },
TokenIdsMismatch {
logprobs: LogprobsCount,
num_token_ids: usize,
},
}
/// Validate logprobs count sampling parameters.
pub(super) fn validate_logprobs(
logprobs: Option<i32>,
prompt_logprobs: Option<i32>,
logprobs: Option<LogprobsCount>,
prompt_logprobs: Option<LogprobsCount>,
logprob_token_ids: Option<&[u32]>,
sampling_limits: SamplingLimits,
) -> Result<(), LogprobsError> {
let vocab_size = sampling_limits.model_vocab_size;
let max_logprobs =
normalize_logprobs_count(sampling_limits.max_logprobs, vocab_size, "max_logprobs")?;
let max_logprobs = sampling_limits.max_logprobs.expanded(vocab_size);
validate_logprobs_count(logprobs, max_logprobs, vocab_size, "logprobs")?;
validate_logprobs_count(prompt_logprobs, max_logprobs, vocab_size, "prompt_logprobs")?;
@@ -50,7 +51,7 @@ pub(super) fn validate_logprobs(
}
fn validate_logprobs_count(
requested: Option<i32>,
requested: Option<LogprobsCount>,
max_logprobs: usize,
vocab_size: usize,
parameter: &'static str,
@@ -59,7 +60,7 @@ fn validate_logprobs_count(
return Ok(());
};
let requested = normalize_logprobs_count(requested, vocab_size, parameter)?;
let requested = requested.expanded(vocab_size);
if requested > max_logprobs {
return Err(LogprobsError::TooManyCount {
parameter,
@@ -72,7 +73,7 @@ fn validate_logprobs_count(
}
pub(super) fn validate_logprob_token_ids(
logprobs: Option<i32>,
logprobs: Option<LogprobsCount>,
logprob_token_ids: Option<&[u32]>,
) -> Result<(), LogprobsError> {
let Some(logprob_token_ids) = logprob_token_ids else {
@@ -88,7 +89,7 @@ pub(super) fn validate_logprob_token_ids(
}
if let Some(logprobs) = logprobs
&& logprobs != n as i32
&& logprobs != LogprobsCount::Top(n as u32)
{
return Err(LogprobsError::TokenIdsMismatch {
logprobs,
@@ -98,15 +99,3 @@ pub(super) fn validate_logprob_token_ids(
Ok(())
}
fn normalize_logprobs_count(
value: i32,
vocab_size: usize,
parameter: &'static str,
) -> Result<usize, LogprobsError> {
match value {
-1 => Ok(vocab_size),
value if value < 0 => Err(LogprobsError::InvalidCount { parameter, value }),
value => Ok(value as usize),
}
}
+8 -1
View File
@@ -309,7 +309,7 @@ fn matches_stop_string(stops: &[String], output: &str, new_bytes: usize) -> Opti
.find_map(|(ss_idx, (ss, len, start_off))| {
output[start_off..]
.windows(len)
.rposition(|w| w == ss)
.position(|w| w == ss)
.map(|pos| (ss_idx, start_off + pos))
})
}
@@ -562,6 +562,13 @@ mod tests {
assert_eq!(result, Some((0, 4)));
}
#[test]
fn stop_string_matches_leftmost_with_multiple_new_bytes() {
let stops = vec!["\n".to_string()];
let result = matches_stop_string(&stops, "Answer\n\n", 2);
assert_eq!(result, Some((0, 6)));
}
#[test]
fn stop_string_matches_at_beginning() {
let stops = vec!["say".to_string()];
+12 -5
View File
@@ -3,9 +3,9 @@ use std::collections::HashMap;
use enum_as_inner::EnumAsInner;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use vllm_engine_core_client::protocol::StructuredOutputsParams;
use vllm_engine_core_client::protocol::lora::LoraRequest;
use vllm_engine_core_client::protocol::multimodal::MmFeatures;
use vllm_engine_core_client::protocol::{LogprobsCount, StructuredOutputsParams};
use crate::error::{Error, Result};
use crate::output::TextDecodeOptions;
@@ -56,14 +56,20 @@ pub struct SamplingParams {
pub max_tokens: Option<u32>,
/// Minimum number of tokens to generate before EOS or stop-token handling.
pub min_tokens: Option<u32>,
/// Maximum number of reasoning ("thinking") tokens to emit before the
/// reasoning section is force-closed. `None` or the user-facing `-1`
/// "unlimited" sentinel both disable the budget. The raw value is carried
/// here; `-1` is normalized to `None` (and other negatives rejected) during
/// lowering (see `lower_sampling_params`).
pub thinking_token_budget: Option<i64>,
/// Number of log probabilities to return per generated token.
///
/// `None` disables sample logprobs. `-1` requests the full vocabulary.
pub logprobs: Option<i32>,
/// `None` disables sample logprobs.
pub logprobs: Option<LogprobsCount>,
/// Number of log probabilities to return per prompt token.
///
/// `None` disables prompt logprobs. `-1` requests the full vocabulary.
pub prompt_logprobs: Option<i32>,
/// `None` disables prompt logprobs.
pub prompt_logprobs: Option<LogprobsCount>,
/// Minimum probability threshold for token sampling. `None` means no
/// explicit user override.
pub min_p: Option<f32>,
@@ -116,6 +122,7 @@ impl Default for SamplingParams {
seed: None,
max_tokens: None,
min_tokens: None,
thinking_token_budget: None,
logprobs: None,
prompt_logprobs: None,
min_p: None,
@@ -15,7 +15,7 @@ from vllm.entrypoints.serve.tokenize.protocol import (
TokenizeChatRequest,
TokenizeCompletionRequest,
)
from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
from vllm.v1.engine.async_llm import AsyncLLM
MODEL_NAME = "openai-community/gpt2"
@@ -58,7 +58,7 @@ class MockModelConfig:
return self.diff_sampling_param or {}
def _build_serving_tokenization(engine: AsyncLLM) -> OpenAIServingTokenization:
def _build_serving_tokenization(engine: AsyncLLM) -> ServingTokenization:
models = OpenAIServingModels(
engine_client=engine,
base_model_paths=BASE_MODEL_PATHS,
@@ -71,8 +71,7 @@ def _build_serving_tokenization(engine: AsyncLLM) -> OpenAIServingTokenization:
chat_template=None,
chat_template_content_format="auto",
)
return OpenAIServingTokenization(
engine,
return ServingTokenization(
models,
openai_serving_render=serving_render,
request_logger=None,
@@ -435,7 +435,7 @@ def test_per_head_quant_scales_backend_selection(
]
+ (
[
("FLASHINFER", True, False), # FlashInfer does not support non-causal
("FLASHINFER", True, True), # FlashInfer supports non-causal
("FLASHINFER", False, True), # FlashInfer works with causal
]
if CudaPlatform is not None
+11 -11
View File
@@ -8,8 +8,8 @@ from vllm.model_executor.kernels.mhc.tilelang import (
_tilelang_hc_prenorm_gemm,
_torch_hc_prenorm_gemm,
)
from vllm.model_executor.layers.mhc import HAS_TILELANG_MHC
from vllm.platforms import current_platform
from vllm.utils.import_utils import has_tilelang
from vllm.utils.torch_utils import set_random_seed
DEVICE = current_platform.device_type
@@ -97,8 +97,8 @@ def hc_head_ref(
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() and has_tilelang()),
reason="CUDA or ROCm and tilelang required",
not HAS_TILELANG_MHC,
reason="TileLang MHC support required",
)
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
@pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -150,8 +150,8 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() and has_tilelang()),
reason="CUDA or ROCm and tilelang required",
not HAS_TILELANG_MHC,
reason="TileLang MHC support required",
)
@pytest.mark.parametrize(
("num_tokens", "hidden_size"),
@@ -190,8 +190,8 @@ def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size):
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() and has_tilelang()),
reason="CUDA or ROCm and tilelang required",
not HAS_TILELANG_MHC,
reason="TileLang MHC support required",
)
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
@pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -217,8 +217,8 @@ def test_mhc_post_tilelang(num_tokens, hidden_size, hc_mult):
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() and has_tilelang()),
reason="CUDA or ROCm and tilelang required",
not HAS_TILELANG_MHC,
reason="TileLang MHC support required",
)
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
@pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -324,8 +324,8 @@ def test_hc_head_triton(num_tokens, hidden_size, hc_mult):
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() and has_tilelang()),
reason="CUDA or ROCm and tilelang required",
not HAS_TILELANG_MHC,
reason="TileLang MHC support required",
)
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
@pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -810,29 +810,6 @@ VLM_TEST_SETTINGS = {
hf_output_post_proc=model_utils.minicpmv_trunc_hf_output,
patch_hf_runner=model_utils.minicpmv_26_patch_hf_runner,
),
"minimax_vl_01": VLMTestInfo(
models=["MiniMaxAI/MiniMax-VL-01"],
prompt_formatter=lambda img_prompt: f"<beginning_of_sentence>user: {img_prompt} assistant:<end_of_sentence>", # noqa: E501
img_idx_to_prompt=lambda _: "<image>",
test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
max_model_len=8192,
max_num_seqs=4,
dtype="bfloat16",
hf_output_post_proc=model_utils.minimax_vl_01_hf_output,
patch_hf_runner=model_utils.minimax_vl_01_patch_hf_runner,
auto_cls=AutoModelForImageTextToText,
marks=[
large_gpu_mark(min_gb=80),
# TODO: [ROCm] Fix pickle issue with ROCm spawn and tp>1
pytest.mark.skipif(
current_platform.is_rocm(),
reason=(
"ROCm: Model too large for single GPU; "
"multi-GPU blocked by HF _LazyConfigMapping pickle issue with spawn"
),
),
],
),
"molmo": VLMTestInfo(
models=["allenai/Molmo-7B-D-0924"],
test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
@@ -245,13 +245,6 @@ def minicpmv_trunc_hf_output(hf_output: RunnerOutput, model: str) -> RunnerOutpu
return output_ids, output_str, out_logprobs
def minimax_vl_01_hf_output(hf_output: RunnerOutput, model: str) -> RunnerOutput:
output_ids, output_str, out_logprobs = hf_output
if output_str.endswith("<end_of_sentence>"):
output_str = output_str.split("<end_of_sentence>")[0]
return output_ids, output_str, out_logprobs
def ultravox_trunc_hf_output(hf_output: RunnerOutput, model: str) -> RunnerOutput:
output_ids, output_str, out_logprobs = hf_output
@@ -1023,17 +1016,6 @@ def minicpmv_26_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
return hf_model
def minimax_vl_01_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
orig_generate = hf_model.model.generate
def _generate(self, *args, image_sizes=None, **kwargs):
return orig_generate(*args, decode_text=False, **kwargs)
hf_model.model.generate = types.MethodType(_generate, hf_model.model)
return hf_model
def molmo_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
"""Patches and returns an instance of the HfRunner to use for Molmo."""
hf_processor = hf_model.processor
@@ -152,3 +152,21 @@ def test_colqwen3_5_relevance_ordering(
dtype: str,
) -> None:
_run_relevance_test(vllm_runner, model, dtype=dtype)
def test_colqwen3_5_config_enables_bidirectional_attention() -> None:
"""ColQwen3.5 retrieval must be served BIDIRECTIONAL (is_causal=False) so the
full_attention layers build with AttentionType.ENCODER_ONLY. This guards the
silent-causal regression (no GPU / model load needed)."""
from types import SimpleNamespace
from vllm.model_executor.models.config import (
MODELS_CONFIG_MAP,
ColQwen3_5Config,
)
assert MODELS_CONFIG_MAP["ColQwen3_5"] is ColQwen3_5Config
model_config = SimpleNamespace(hf_config=SimpleNamespace())
ColQwen3_5Config.verify_and_update_model_config(model_config)
assert model_config.hf_config.is_causal is False
@@ -1,113 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from PIL import Image
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.parse import ImageSize
from vllm.multimodal.processing import BaseMultiModalProcessor
from ....conftest import ImageTestAssets
from ...utils import build_model_context
@pytest.mark.parametrize("model_id", ["MiniMaxAI/MiniMax-VL-01"])
@pytest.mark.parametrize("num_imgs", [1, 2])
def test_processor_override(
image_assets: ImageTestAssets,
model_id: str,
num_imgs: int,
):
ctx = build_model_context(
model_id,
mm_processor_kwargs=None,
limit_mm_per_prompt={"image": num_imgs},
)
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
prompt = "<image>" * num_imgs
image = Image.new("RGB", size=(364, 364))
mm_data = {"image": [image] * num_imgs}
processed_inputs = processor(
prompt,
mm_items=processor.info.parse_mm_data(mm_data),
hf_processor_mm_kwargs={},
)
image_placeholders = processed_inputs["mm_placeholders"]["image"]
assert len(image_placeholders) == num_imgs
def _validate_image_prompt_replacements_one(
processor: BaseMultiModalProcessor,
num_imgs: int,
failed_size_excs: list[tuple[ImageSize, Exception]],
image_size: ImageSize,
) -> None:
prompt = "<image>" * num_imgs
image = Image.new("RGB", size=image_size)
mm_data = {"image": [image] * num_imgs}
try:
processed_inputs = processor(
prompt,
mm_items=processor.info.parse_mm_data(mm_data),
hf_processor_mm_kwargs={},
)
image_placeholders = processed_inputs["mm_placeholders"]["image"]
assert len(image_placeholders) == num_imgs
except Exception as exc:
failed_size_excs.append((image_size, exc))
def _test_image_prompt_replacements(
processor,
*,
num_imgs: int,
image_sizes: list[ImageSize],
) -> None:
failed_size_excs = list[tuple[ImageSize, Exception]]()
for size in image_sizes:
_validate_image_prompt_replacements_one(
processor, num_imgs, failed_size_excs, size
)
if failed_size_excs:
msg = "Found failing image sizes:" + "\n========\n".join(
f"[{size}]\n{exc}" for size, exc in failed_size_excs
)
raise AssertionError(msg)
@pytest.mark.parametrize("model_id", ["MiniMaxAI/MiniMax-VL-01"])
@pytest.mark.parametrize("num_imgs", [1, 2])
def test_processor_prompt_replacements_regression(model_id, num_imgs):
ctx = build_model_context(
model_id,
mm_processor_kwargs=None,
limit_mm_per_prompt={"image": num_imgs},
)
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
image_ratios = [
(171, 152),
(184, 161),
(198, 176),
(333, 296),
(369, 328),
(488, 183),
(2560, 1669),
]
image_sizes = [
size for w, h in image_ratios for size in [ImageSize(w, h), ImageSize(h, w)]
]
_test_image_prompt_replacements(
processor,
num_imgs=num_imgs,
image_sizes=image_sizes,
)
-13
View File
@@ -421,15 +421,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
},
trust_remote_code=True,
),
"MiniMaxForCausalLM": _HfExamplesInfo("MiniMaxAI/MiniMax-Text-01-hf"),
"MiniMaxText01ForCausalLM": _HfExamplesInfo(
"MiniMaxAI/MiniMax-Text-01",
trust_remote_code=True,
revision="a59aa9cbc53b9fb8742ca4e9e1531b9802b6fdc3",
),
"MiniMaxM1ForCausalLM": _HfExamplesInfo(
"MiniMaxAI/MiniMax-M1-40k", trust_remote_code=True
),
"MiniMaxM2ForCausalLM": _HfExamplesInfo(
"MiniMaxAI/MiniMax-M2",
trust_remote_code=True,
@@ -1113,10 +1104,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"openbmb/MiniCPM-V-4_6",
min_transformers_version="5.7.0",
),
"MiniMaxVL01ForConditionalGeneration": _HfExamplesInfo(
"MiniMaxAI/MiniMax-VL-01",
trust_remote_code=True,
),
"MiniMaxM3SparseForConditionalGeneration": _HfExamplesInfo(
"MiniMaxAI/MiniMax-M3",
trust_remote_code=True,
-5
View File
@@ -98,11 +98,6 @@ def can_initialize(
vllm_config.validate_block_size()
return scheduler_kv_cache_config
if model_arch == "MiniMaxVL01ForConditionalGeneration":
pytest.skip(
"pickle error when loading `transformers.models.auto.CONFIG_MAPPING`"
)
if model_arch == "MoonshotKimiaForCausalLM":
pytest.skip(
"Kimi-Audio requires SpeechToTextConfig "
+7 -1
View File
@@ -507,7 +507,13 @@ def dummy_hf_overrides(
# Only set MoE related config when the model has MoE layers.
# Otherwise all models detected as MoE by _get_transformers_backend_cls.
if model_arch_config.num_experts > 0:
num_experts_per_tok = 1 if model_arch == "Llama4ForConditionalGeneration" else 2
num_experts_per_tok = 2
if model_arch in (
"Llama4ForConditionalGeneration",
"Llama4ForCausalLM",
"EagleLlama4ForCausalLM",
):
num_experts_per_tok = 1
update_dict.update(
{
"num_experts": num_experts,
File diff suppressed because it is too large Load Diff
@@ -6,10 +6,12 @@ from types import SimpleNamespace
import pytest
from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import (
MiniMaxText01LinearAttention,
)
from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
from vllm.model_executor.layers.mamba.short_conv import ShortConv
from vllm.model_executor.models.minimax_text_01 import MiniMaxText01LinearAttention
from vllm.v1.attention.backends.linear_attn import LinearAttentionBackend
from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionBackend
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend
@@ -16,6 +16,7 @@ from tests.v1.attention.utils import (
create_vllm_config,
)
from vllm.config import SpeculativeConfig
from vllm.config.compilation import CUDAGraphMode
from vllm.v1.attention.backends.gdn_attn import (
GDNAttentionMetadata,
GDNAttentionMetadataBuilder,
@@ -123,9 +124,15 @@ GDN_BUILD_TEST_CASES = {
def _create_gdn_builder(
num_speculative_tokens: int = 0,
full_cuda_graph: bool = False,
) -> GDNAttentionMetadataBuilder:
"""Create a GDNAttentionMetadataBuilder with minimal config."""
vllm_config = create_vllm_config(block_size=BLOCK_SIZE)
vllm_config = create_vllm_config(
model_name="Qwen/Qwen3.5-0.8B",
block_size=BLOCK_SIZE,
)
if full_cuda_graph:
vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
if num_speculative_tokens > 0:
vllm_config.speculative_config = SpeculativeConfig(
method="ngram",
@@ -189,3 +196,28 @@ def test_has_initial_state_after_reclassification():
assert meta.has_initial_state is not None
# req0 has context_lens = 65 - 1 = 64 > 0, so has_initial_state[0] = True
assert meta.has_initial_state[0].item() is True
def test_full_cudagraph_spec_metadata_uses_request_count():
"""FULL cudagraph token padding must not pad request-indexed metadata."""
num_speculative_tokens = 3
builder = _create_gdn_builder(
num_speculative_tokens=num_speculative_tokens,
full_cuda_graph=True,
)
batch = BatchSpec(seq_lens=[80, 96], query_lens=[4, 4])
meta = _build(builder, batch, num_decode_draft_tokens=[3, 3])
assert meta.num_spec_decodes == batch.batch_size
assert meta.num_spec_decode_tokens == batch.compute_num_tokens()
assert meta.spec_state_indices_tensor is not None
assert meta.spec_state_indices_tensor.shape == (
batch.batch_size,
num_speculative_tokens + 1,
)
assert meta.spec_sequence_masks is not None
assert meta.spec_sequence_masks.shape == (batch.batch_size,)
assert meta.spec_query_start_loc is not None
assert meta.spec_query_start_loc.shape == (batch.batch_size + 1,)
assert meta.num_accepted_tokens is not None
assert meta.num_accepted_tokens.shape == (batch.batch_size,)
@@ -0,0 +1,460 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Callable
import pytest
import vllm.v1.core.kv_cache_utils as kv_cache_utils
from vllm.distributed.kv_events import BlockRemoved, BlockStored
from vllm.sampling_params import SamplingParams
from vllm.utils.hashing import sha256
from vllm.v1.core.block_pool import BlockPool
from vllm.v1.core.kv_cache_utils import (
BlockHash,
BlockHashListWithBlockSize,
KVCacheBlock,
get_request_block_hasher,
hash_block_tokens,
init_none_hash,
)
from vllm.v1.request import Request
pytestmark = pytest.mark.cpu_test
@pytest.fixture(autouse=True)
def _auto_init_hash_fn():
init_none_hash(sha256)
def make_request(
request_id: str,
prompt_token_ids: list[int],
hash_block_size: int,
hash_fn: Callable,
) -> Request:
sampling_params = SamplingParams(max_tokens=17)
sampling_params.update_from_generation_config({}, eos_token_id=100)
return Request(
request_id=request_id,
prompt_token_ids=prompt_token_ids,
sampling_params=sampling_params,
pooling_params=None,
block_hasher=get_request_block_hasher(hash_block_size, hash_fn),
)
def boundary_hash(req: Request, hash_block_size: int, num_tokens: int) -> BlockHash:
# Every boundary at a hash_block_size multiple is just the fine-grained
# chain hash ending there.
return req.block_hashes[num_tokens // hash_block_size - 1]
def cache_full_block_and_partial_tail(
token_ids: list[int],
*,
enable_kv_cache_events: bool = False,
) -> tuple[BlockPool, Request, list[KVCacheBlock], BlockHash]:
hash_block_size = 2
block_size = 6
kv_cache_group_id = 0
req = make_request("0", token_ids, hash_block_size, sha256)
pool = BlockPool(
num_gpu_blocks=3,
enable_caching=True,
hash_block_size=hash_block_size,
enable_kv_cache_events=enable_kv_cache_events,
)
blocks = pool.get_new_blocks(2)
pool.cache_full_blocks(
request=req,
blocks=blocks,
num_cached_blocks=0,
num_full_blocks=1,
block_size=block_size,
kv_cache_group_id=kv_cache_group_id,
)
partial_hash = boundary_hash(req, hash_block_size, len(token_ids))
assert pool.cache_partial_block(
request=req,
block=blocks[1],
num_tokens=len(token_ids),
kv_cache_group_id=kv_cache_group_id,
block_size=block_size,
)
return pool, req, blocks, partial_hash
def test_boundary_hashes_reuse_fine_grained_chain():
hash_block_size = 2
block_size = 6
token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
req = make_request("0", token_ids, hash_block_size, sha256)
coarse = BlockHashListWithBlockSize(req.block_hashes, hash_block_size, block_size)
# The block_size=6 full-block hash is the fine hash at the 6-token boundary,
# not a concatenation of the three fine hashes inside the block.
assert coarse[0] == req.block_hashes[6 // hash_block_size - 1]
assert coarse[0] != BlockHash(
req.block_hashes[0] + req.block_hashes[1] + req.block_hashes[2]
)
# A partial tail at 10 tokens is the fine hash at the 10-token boundary,
# which chains over the entire prefix.
tail_hash = boundary_hash(req, hash_block_size, 10)
assert tail_hash == req.block_hashes[4]
assert tail_hash == hash_block_tokens(sha256, req.block_hashes[3], token_ids[8:10])
def test_cache_partial_block_kv_cache_events():
hash_block_size = 4
block_size = 12
kv_cache_group_id = 2
pool = BlockPool(
num_gpu_blocks=2,
enable_caching=True,
hash_block_size=hash_block_size,
enable_kv_cache_events=True,
)
req = make_request(
"req_partial_events",
prompt_token_ids=list(range(hash_block_size * 2)),
hash_block_size=hash_block_size,
hash_fn=sha256,
)
block = pool.get_new_blocks(1)[0]
partial_entry_hash = pool.cache_partial_block(
request=req,
block=block,
num_tokens=hash_block_size * 2,
kv_cache_group_id=kv_cache_group_id,
block_size=block_size,
)
events = pool.take_events()
assert len(events) == 1
stored_event = events[0]
assert isinstance(stored_event, BlockStored)
assert partial_entry_hash is not None
assert stored_event.block_hashes == [
kv_cache_utils.maybe_convert_block_hash(req.block_hashes[1])
]
assert stored_event.parent_block_hash == kv_cache_utils.maybe_convert_block_hash(
req.block_hashes[0]
)
assert stored_event.token_ids == req.all_token_ids[hash_block_size:]
assert stored_event.block_size == 4
assert stored_event.group_idx == kv_cache_group_id
duplicate_entry_hash = pool.cache_partial_block(
request=req,
block=block,
num_tokens=hash_block_size * 2,
kv_cache_group_id=kv_cache_group_id,
block_size=block_size,
)
assert duplicate_entry_hash == partial_entry_hash
assert pool.take_events() == []
pool.free_blocks([block])
pool.get_new_blocks(1)
events = pool.take_events()
assert len(events) == 1
removed_event = events[0]
assert isinstance(removed_event, BlockRemoved)
assert removed_event.block_hashes == stored_event.block_hashes
assert removed_event.group_idx == kv_cache_group_id
def test_partial_block_replacement_emits_remove_then_store_events():
hash_block_size = 2
block_size = 6
kv_cache_group_id = 0
req = make_request("0", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)
pool = BlockPool(
num_gpu_blocks=3,
enable_caching=True,
hash_block_size=hash_block_size,
enable_kv_cache_events=True,
)
blocks = pool.get_new_blocks(2)
pool.cache_full_blocks(
request=req,
blocks=blocks,
num_cached_blocks=0,
num_full_blocks=1,
block_size=block_size,
kv_cache_group_id=kv_cache_group_id,
)
partial_hash_8 = boundary_hash(req, hash_block_size, 8)
assert pool.cache_partial_block(
request=req,
block=blocks[1],
num_tokens=8,
kv_cache_group_id=kv_cache_group_id,
block_size=block_size,
)
assert pool.get_cached_block(partial_hash_8, [kv_cache_group_id]) == [blocks[1]]
pool.take_events()
req.append_output_token_ids([4, 4])
partial_hash_10 = boundary_hash(req, hash_block_size, 10)
assert pool.cache_partial_block(
request=req,
block=blocks[1],
num_tokens=10,
kv_cache_group_id=kv_cache_group_id,
block_size=block_size,
)
events = pool.take_events()
assert len(events) == 2
removed_event, stored_event = events
assert isinstance(removed_event, BlockRemoved)
assert removed_event.block_hashes == [
kv_cache_utils.maybe_convert_block_hash(partial_hash_8)
]
assert removed_event.group_idx == kv_cache_group_id
assert isinstance(stored_event, BlockStored)
assert stored_event.block_hashes == [
kv_cache_utils.maybe_convert_block_hash(partial_hash_10)
]
assert stored_event.parent_block_hash == kv_cache_utils.maybe_convert_block_hash(
boundary_hash(req, hash_block_size, 8)
)
assert stored_event.token_ids == req.all_token_ids[8:10]
assert stored_event.block_size == hash_block_size
assert stored_event.group_idx == kv_cache_group_id
assert pool.get_cached_block(partial_hash_8, [kv_cache_group_id]) is None
assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) == [blocks[1]]
def test_later_request_hits_cached_partial_tail():
hash_block_size = 2
block_size = 6
kv_cache_group_id = 0
cached_token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
req = make_request("0", cached_token_ids, hash_block_size, sha256)
pool = BlockPool(
num_gpu_blocks=3,
enable_caching=True,
hash_block_size=hash_block_size,
)
blocks = pool.get_new_blocks(2)
pool.cache_full_blocks(
request=req,
blocks=blocks,
num_cached_blocks=0,
num_full_blocks=1,
block_size=block_size,
kv_cache_group_id=kv_cache_group_id,
)
partial_hash_10 = boundary_hash(req, hash_block_size, 10)
assert pool.cache_partial_block(
request=req,
block=blocks[1],
num_tokens=10,
kv_cache_group_id=kv_cache_group_id,
block_size=block_size,
)
replay = make_request("1", cached_token_ids, hash_block_size, sha256)
replay_hash_10 = boundary_hash(replay, hash_block_size, 10)
assert replay_hash_10 == partial_hash_10
assert pool.get_cached_block(replay_hash_10, [kv_cache_group_id]) == [blocks[1]]
extended = make_request("2", cached_token_ids + [10], hash_block_size, sha256)
extended_hash_10 = boundary_hash(extended, hash_block_size, 10)
assert extended_hash_10 == partial_hash_10
assert pool.get_cached_block(extended_hash_10, [kv_cache_group_id]) == [blocks[1]]
def test_cache_partial_block_uses_fine_grained_boundary_hash():
hash_block_size = 2
block_size = 6
kv_cache_group_id = 0
token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
req = make_request("0", token_ids, hash_block_size, sha256)
pool = BlockPool(
num_gpu_blocks=3,
enable_caching=True,
hash_block_size=hash_block_size,
)
blocks = pool.get_new_blocks(2)
pool.cache_full_blocks(
request=req,
blocks=blocks,
num_cached_blocks=0,
num_full_blocks=1,
block_size=block_size,
kv_cache_group_id=kv_cache_group_id,
)
partial_entry_hash = pool.cache_partial_block(
request=req,
block=blocks[1],
num_tokens=10,
kv_cache_group_id=kv_cache_group_id,
block_size=block_size,
)
# The partial entry is keyed by the fine-grained hash at the 10-token
# boundary, regardless of the owning group's block_size.
expected = boundary_hash(req, hash_block_size, 10)
assert partial_entry_hash == kv_cache_utils.make_block_hash_with_group_id(
expected, kv_cache_group_id
)
assert pool.get_cached_block(expected, [kv_cache_group_id]) == [blocks[1]]
def test_cache_partial_block_requires_hash_boundary():
hash_block_size = 2
block_size = 4
req = make_request("0", [0, 0, 1, 1], hash_block_size, sha256)
pool = BlockPool(
num_gpu_blocks=2,
enable_caching=True,
hash_block_size=hash_block_size,
)
block = pool.get_new_blocks(1)[0]
with pytest.raises(AssertionError):
pool.cache_partial_block(
request=req,
block=block,
num_tokens=3,
kv_cache_group_id=0,
block_size=block_size,
)
def test_cache_partial_block_duplicate_checks_all_blocks_for_hash():
hash_block_size = 2
block_size = 4
kv_cache_group_id = 0
req = make_request("0", [0, 0, 1, 1], hash_block_size, sha256)
pool = BlockPool(
num_gpu_blocks=4,
enable_caching=True,
hash_block_size=hash_block_size,
)
blocks = pool.get_new_blocks(2)
first_entry_hash = pool.cache_partial_block(
request=req,
block=blocks[0],
num_tokens=2,
kv_cache_group_id=kv_cache_group_id,
block_size=block_size,
)
second_entry_hash = pool.cache_partial_block(
request=req,
block=blocks[1],
num_tokens=2,
kv_cache_group_id=kv_cache_group_id,
block_size=block_size,
)
assert first_entry_hash == second_entry_hash
duplicate_entry_hash = pool.cache_partial_block(
request=req,
block=blocks[1],
num_tokens=2,
kv_cache_group_id=kv_cache_group_id,
block_size=block_size,
)
assert duplicate_entry_hash == second_entry_hash
assert pool.cached_block_hashes_by_block == {}
def test_reset_prefix_cache_clears_partial_entry_metadata():
pool, req, blocks, partial_hash_10 = cache_full_block_and_partial_tail(
[0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
)
full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, 6)[0]
assert pool.get_cached_block(full_hash, [0]) == [blocks[0]]
assert pool.get_cached_block(partial_hash_10, [0]) == [blocks[1]]
pool.free_blocks(blocks)
assert pool.reset_prefix_cache()
assert pool.get_cached_block(full_hash, [0]) is None
assert pool.get_cached_block(partial_hash_10, [0]) is None
assert pool.cached_block_hashes_by_block == {}
def test_evict_cached_block_removes_full_hash_and_partial_entry():
pool, req, blocks, partial_hash_10 = cache_full_block_and_partial_tail(
[0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
)
full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, 6)[0]
assert pool.get_cached_block(full_hash, [0]) == [blocks[0]]
assert pool.get_cached_block(partial_hash_10, [0]) == [blocks[1]]
pool.evict_blocks({blocks[0].block_id, blocks[1].block_id})
assert pool.get_cached_block(full_hash, [0]) is None
assert pool.get_cached_block(partial_hash_10, [0]) is None
assert pool.cached_block_hashes_by_block == {}
def test_partial_block_promotes_to_direct_full_block_hash():
hash_block_size = 2
block_size = 6
kv_cache_group_id = 0
token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
req = make_request("0", token_ids, hash_block_size, sha256)
pool = BlockPool(
num_gpu_blocks=3,
enable_caching=True,
hash_block_size=hash_block_size,
)
blocks = pool.get_new_blocks(2)
pool.cache_full_blocks(
request=req,
blocks=blocks,
num_cached_blocks=0,
num_full_blocks=1,
block_size=block_size,
kv_cache_group_id=kv_cache_group_id,
)
partial_hash_10 = boundary_hash(req, hash_block_size, 10)
assert pool.cache_partial_block(
request=req,
block=blocks[1],
num_tokens=10,
kv_cache_group_id=kv_cache_group_id,
block_size=block_size,
)
assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) == [blocks[1]]
req.append_output_token_ids([5, 5])
full_hashes = BlockHashListWithBlockSize(
req.block_hashes, hash_block_size, block_size
)
promoted_full_hash = full_hashes[1]
# The promoted full-block hash is the fine hash at the 12-token boundary,
# not a concatenation of the fine hashes inside the block.
assert promoted_full_hash == req.block_hashes[12 // hash_block_size - 1]
assert promoted_full_hash != BlockHash(
req.block_hashes[3] + req.block_hashes[4] + req.block_hashes[5]
)
pool.cache_full_blocks(
request=req,
blocks=blocks,
num_cached_blocks=1,
num_full_blocks=2,
block_size=block_size,
kv_cache_group_id=kv_cache_group_id,
)
assert pool.get_cached_block(promoted_full_hash, [kv_cache_group_id]) == [blocks[1]]
assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) is None
+1 -1
View File
@@ -225,7 +225,7 @@ def test_kv_cache_block():
# Test block hash setting and resetting
block_hash = make_block_hash_with_group_id(BlockHash(b"abc"), 0)
block.block_hash = block_hash
block.set_block_hash(block_hash)
assert block.block_hash == block_hash
block.reset_hash()
+1 -1
View File
@@ -2003,7 +2003,7 @@ def test_maybe_evict_cached_block():
assert len(pool.blocks) == len(block_hashes)
# Manually add all blocks to cached_blocks
for block, block_hash in zip(pool.blocks, block_hashes):
block.block_hash = block_hash
block.set_block_hash(block_hash)
pool.cached_block_hash_to_block.insert(block_hash, block)
block0, block1, block2, block3 = pool.blocks
@@ -0,0 +1,355 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock
import pytest
import torch
from tests.v1.kv_connector.unit.utils import create_vllm_config
from vllm.config import KVEventsConfig, KVTransferConfig
from vllm.distributed.kv_events import BlockRemoved, BlockStored
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import (
OffloadingEventGroupSpec,
OffloadingEventsTracker,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
GroupOffloadConfig,
)
from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheSpecKind,
)
from vllm.v1.kv_offload.base import (
OffloadingEvent,
OffloadingKVEventsConfig,
OffloadKey,
make_offload_key,
)
from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
_CPU_MEDIUM = CPULoadStoreSpec.medium()
_FULL_ATTENTION_EVENT_SPEC = OffloadingEventGroupSpec(
kv_cache_spec_kind=KVCacheSpecKind.FULL_ATTENTION.value,
kv_cache_spec_sliding_window=None,
)
def _tracker(
*,
enable_kv_cache_events: bool = True,
self_describing_kv_events: bool = True,
) -> OffloadingEventsTracker:
return OffloadingEventsTracker(
OffloadingKVEventsConfig(
enable_kv_cache_events=enable_kv_cache_events,
self_describing_kv_events=self_describing_kv_events,
)
)
def _hash(i: int) -> BlockHash:
return BlockHash(str(i).encode())
def _wire_hash(block_hash: BlockHash):
return maybe_convert_block_hash(block_hash)
def _request(*, block_hashes: list[BlockHash], token_count: int):
req = MagicMock()
req.block_hashes = block_hashes
req.all_token_ids = list(range(1, token_count + 1))
req.lora_request = None
return req
def _group_config(
*,
group_idx: int = 0,
block_size: int = 4,
block_size_factor: int = 1,
sliding_window_size_in_blocks: int | None = None,
) -> GroupOffloadConfig:
return GroupOffloadConfig(
group_idx=group_idx,
gpu_block_size=block_size,
offloaded_block_size=block_size * block_size_factor,
hash_block_size_factor=block_size_factor,
sliding_window_size_in_blocks=sliding_window_size_in_blocks,
kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC,
)
def _record_chunks(
tracker: OffloadingEventsTracker,
req,
group_config: GroupOffloadConfig,
num_chunks: int,
) -> list[OffloadKey]:
keys: list[OffloadKey] = []
hbf = group_config.hash_block_size_factor
for chunk_idx in range(num_chunks):
tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1]
assert tail_hash is not None
key = make_offload_key(tail_hash, group_config.group_idx)
tracker.record_store(req, group_config, chunk_idx, key)
keys.append(key)
return keys
def _stored_event(keys: list[OffloadKey]) -> OffloadingEvent:
return OffloadingEvent(keys=keys, medium=_CPU_MEDIUM, removed=False)
def _removed_event(keys: list[OffloadKey]) -> OffloadingEvent:
return OffloadingEvent(keys=keys, medium=_CPU_MEDIUM, removed=True)
def test_take_events_publishes_routable_block_stored():
block_size = 4
tracker = _tracker()
group_config = _group_config(block_size=block_size)
req = _request(
block_hashes=[_hash(i) for i in range(6)],
token_count=block_size * 6,
)
keys = _record_chunks(tracker, req, group_config, num_chunks=6)
batch1 = list(tracker.take_events([_stored_event(keys[:3])]))
assert len(batch1) == 3
for i, event in enumerate(batch1):
assert isinstance(event, BlockStored)
assert event.medium == _CPU_MEDIUM
assert event.block_hashes == [_wire_hash(_hash(i))]
assert event.block_size == block_size
assert event.token_ids == list(
range(i * block_size + 1, (i + 1) * block_size + 1)
)
if i == 0:
assert event.parent_block_hash is None
else:
assert event.parent_block_hash == _wire_hash(_hash(i - 1))
assert event.lora_id is None
assert event.lora_name is None
assert event.extra_keys is None
assert event.group_idx == 0
assert event.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
assert event.kv_cache_spec_sliding_window is None
batch2 = list(tracker.take_events([_stored_event(keys[3:])]))
assert len(batch2) == 3
assert batch2[0].parent_block_hash == batch1[-1].block_hashes[-1]
assert len(tracker._pending_event_metadata) == 6
def test_take_events_factor_gt_1_chunk_store_and_remove():
block_size = 4
block_size_factor = 3
tracker = _tracker()
group_config = _group_config(
block_size=block_size, block_size_factor=block_size_factor
)
req = _request(
block_hashes=[_hash(i) for i in range(6)],
token_count=block_size * block_size_factor * 2,
)
keys = _record_chunks(tracker, req, group_config, num_chunks=2)
stored = list(tracker.take_events([_stored_event(keys)]))
assert len(stored) == 2
expected_hashes = []
for chunk_idx, event in enumerate(stored):
assert isinstance(event, BlockStored)
expected_chunk_hashes = [
_wire_hash(_hash(i))
for i in range(
chunk_idx * block_size_factor,
(chunk_idx + 1) * block_size_factor,
)
]
assert event.block_hashes == expected_chunk_hashes
assert event.block_size == block_size
assert len(event.token_ids) == block_size * block_size_factor
if chunk_idx == 0:
assert event.parent_block_hash is None
else:
assert event.parent_block_hash == _wire_hash(_hash(block_size_factor - 1))
expected_hashes.extend(expected_chunk_hashes)
assert len(tracker._pending_event_metadata) == 2
removed = list(tracker.take_events([_removed_event(keys)]))
assert len(removed) == 1
assert isinstance(removed[0], BlockRemoved)
assert removed[0].block_hashes == expected_hashes
assert removed[0].medium == _CPU_MEDIUM
assert removed[0].group_idx == 0
assert not tracker._pending_event_metadata
def test_take_events_factor_gt_1_store_is_order_independent():
block_size_factor = 3
tracker = _tracker()
group_config = _group_config(block_size_factor=block_size_factor)
req = _request(
block_hashes=[_hash(i) for i in range(6)],
token_count=4 * block_size_factor * 2,
)
keys = _record_chunks(tracker, req, group_config, num_chunks=2)
unknown_key = make_offload_key(_hash(12345), 0)
events = list(tracker.take_events([_stored_event([keys[1], unknown_key, keys[0]])]))
assert len(events) == 3
chunk1, placeholder, chunk0 = events
assert [len(event.block_hashes) for event in events] == [3, 1, 3]
assert placeholder.block_size == 0
assert placeholder.token_ids == []
assert chunk0.parent_block_hash is None
assert chunk1.parent_block_hash == chunk0.block_hashes[-1]
def test_take_events_opt_out_keeps_placeholders():
tracker = _tracker(self_describing_kv_events=False)
group_config = _group_config()
req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
keys = _record_chunks(tracker, req, group_config, num_chunks=3)
assert not tracker.self_describing_enabled
assert not tracker._pending_event_metadata
events = list(
tracker.take_events(
[
_stored_event(keys),
_removed_event(keys),
]
)
)
assert len(events) == 4
for event in events[:3]:
assert isinstance(event, BlockStored)
assert event.block_size == 0
assert event.token_ids == []
assert event.parent_block_hash is None
assert isinstance(events[3], BlockRemoved)
assert len(events[3].block_hashes) == 3
def test_record_store_skips_sliding_window_group():
tracker = _tracker()
group_config = _group_config(sliding_window_size_in_blocks=2)
req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
keys = _record_chunks(tracker, req, group_config, num_chunks=3)
assert not tracker._pending_event_metadata
events = list(tracker.take_events([_stored_event(keys[:1])]))
assert len(events) == 1
assert isinstance(events[0], BlockStored)
assert events[0].block_size == 0
def test_take_events_groups_removed_hashes_by_kv_group():
tracker = _tracker()
group0_config = _group_config(group_idx=0, block_size_factor=2)
group1_config = _group_config(group_idx=1, block_size_factor=2)
req0 = _request(block_hashes=[_hash(0), _hash(1)], token_count=8)
req1 = _request(block_hashes=[_hash(10), _hash(11)], token_count=8)
key0 = _record_chunks(tracker, req0, group0_config, num_chunks=1)[0]
key1 = _record_chunks(tracker, req1, group1_config, num_chunks=1)[0]
removed = list(tracker.take_events([_removed_event([key0, key1])]))
assert len(removed) == 2
by_group = {event.group_idx: event.block_hashes for event in removed}
assert by_group == {
0: [_wire_hash(_hash(0)), _wire_hash(_hash(1))],
1: [_wire_hash(_hash(10)), _wire_hash(_hash(11))],
}
def test_take_events_supports_restore_after_eviction():
block_size = 4
tracker = _tracker()
group_config = _group_config(block_size=block_size)
req = _request(block_hashes=[_hash(0)], token_count=block_size)
key = _record_chunks(tracker, req, group_config, num_chunks=1)[0]
first_store = list(tracker.take_events([_stored_event([key])]))
assert len(first_store) == 1
assert isinstance(first_store[0], BlockStored)
assert first_store[0].token_ids == [1, 2, 3, 4]
removed = list(tracker.take_events([_removed_event([key])]))
assert len(removed) == 1
assert isinstance(removed[0], BlockRemoved)
assert not tracker._pending_event_metadata
req.all_token_ids = [5, 6, 7, 8]
tracker.record_store(req, group_config, offload_block_idx=0, offload_key=key)
second_store = list(tracker.take_events([_stored_event([key])]))
assert len(second_store) == 1
assert isinstance(second_store[0], BlockStored)
assert second_store[0].token_ids == [5, 6, 7, 8]
def test_reset_cache_clears_side_table():
tracker = _tracker()
group_config = _group_config()
req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
_record_chunks(tracker, req, group_config, num_chunks=3)
assert tracker._pending_event_metadata
tracker.reset()
assert not tracker._pending_event_metadata
def test_tiering_rejects_self_describing_kv_events():
vllm_config = create_vllm_config(
block_size=4,
max_num_batched_tokens=16,
disable_hybrid_kv_cache_manager=False,
)
vllm_config.kv_transfer_config = KVTransferConfig(
kv_connector="OffloadingConnector",
kv_role="kv_both",
kv_connector_extra_config={
"spec_name": "TieringOffloadingSpec",
"cpu_bytes_to_use": 1 << 20,
"self_describing_kv_events": True,
"secondary_tiers": [{"type": "example"}],
},
)
vllm_config.kv_events_config = KVEventsConfig(
enable_kv_cache_events=True,
publisher="null",
)
kv_cache_config = KVCacheConfig(
num_blocks=0,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(
["layer"],
FullAttentionSpec(
block_size=4,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
)
],
)
with pytest.raises(ValueError, match="TieringOffloadingSpec"):
TieringOffloadingSpec(vllm_config, kv_cache_config)
@@ -1,6 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable
from types import SimpleNamespace
from unittest.mock import MagicMock
@@ -12,19 +11,16 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
to_keys,
)
from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
from vllm.distributed.kv_events import BlockRemoved, BlockStored
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
OffloadingConnectorScheduler,
RequestOffloadState,
)
from vllm.v1.core.kv_cache_utils import BlockHash
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheGroupSpec,
SlidingWindowSpec,
)
from vllm.v1.kv_offload.base import (
OffloadingEvent,
OffloadingManager,
OffloadPolicy,
ReqContext,
@@ -146,31 +142,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5))
# test take_events
def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
return [BlockHash(str(i).encode()) for i in int_hashes]
def take_events() -> Iterable[OffloadingEvent]:
yield OffloadingEvent(keys=to_keys([1, 2, 3]), medium="A", removed=False)
yield OffloadingEvent(keys=to_keys([4, 5, 6]), medium="B", removed=True)
runner.manager.take_events.side_effect = take_events
events = list(runner.scheduler_connector.take_events())
assert len(events) == 2
event = events[0]
assert isinstance(event, BlockStored)
assert event.block_hashes == to_hashes([1, 2, 3])
assert event.block_size == 0
assert event.medium == "A"
assert event.token_ids == []
assert event.parent_block_hash is None
assert event.lora_id is None
assert event.lora_name is None
event = events[1]
assert isinstance(event, BlockRemoved)
assert event.block_hashes == to_hashes([4, 5, 6])
assert event.medium == "B"
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_request_preemption(request_runner, async_scheduling: bool):
@@ -14,7 +14,12 @@ from tests.v1.kv_connector.unit.utils import (
create_vllm_config,
)
from vllm import SamplingParams
from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config
from vllm.config import (
KVEventsConfig,
KVTransferConfig,
VllmConfig,
set_current_vllm_config,
)
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
OffloadingConnectorMetadata,
@@ -198,6 +203,9 @@ class RequestRunner:
"spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils", # noqa: E501
# Preserve legacy behavior for tests; new opt-in tests override.
"offload_prompt_only": False,
# Exercise the self-describing KV events path by default;
# opt-out tests override this to cover the legacy placeholders.
"self_describing_kv_events": True,
}
if block_size_factor > 1:
extra_config["block_size"] = block_size * block_size_factor
@@ -209,6 +217,13 @@ class RequestRunner:
kv_role="kv_both",
kv_connector_extra_config=extra_config,
)
vllm_config.kv_events_config = KVEventsConfig(
# Enable so the offloading events tracker is active, but use the
# null publisher: these tests drain take_events directly and a
# real ZMQ publisher would bind a port per test.
enable_kv_cache_events=True,
publisher="null",
)
if kv_cache_groups is None:
kv_cache_groups = [
@@ -219,6 +219,7 @@ def test_multi_example_connector_consistency():
enforce_eager=True,
gpu_memory_utilization=0.5,
kv_transfer_config=kv_transfer_config,
async_scheduling=False,
)
# Run generation - this should trigger saving KV cache
# Use a single prompt to avoid race conditions depending on the order of scheduling
@@ -138,6 +138,10 @@ def _wait_for_prefix_cache_reset(llm: LLM) -> None:
def _latency_test(llm: LLM, subscriber: MockSubscriber | None):
# TODO: Reintroduce latency test on ROCm once MRV2 supports cross
# layer KV Cache. See https://github.com/vllm-project/vllm/pull/45947
if current_platform.is_rocm():
return
sampling_params = SamplingParams(max_tokens=1)
num_times_cpu_better_than_cold = 0
+8 -8
View File
@@ -294,25 +294,25 @@ def test_cpu_manager():
# prepare store with no space ([2, 3] is being loaded)
assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None
# complete load [2, 3]
# complete load [2, 3]. Load changes the eviction list, making 2, 3 recent.
cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
# prepare store [6, 7, 8] -> evicts [2, 3, 4] (oldest)
# prepare store [6, 7, 8] -> evicts [4, 5, 2] (oldest)
prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
verify_store_output(
prepare_store_output,
ExpectedPrepareStoreOutput(
keys_to_store=[6, 7, 8],
store_block_ids=[3, 2, 1],
evicted_keys=[2, 3, 4],
store_block_ids=[1, 0, 3],
evicted_keys=[4, 5, 2],
),
)
# complete store [6, 7, 8]
cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
# touch [5, 6, 7] (move to end of LRU order)
cpu_manager.touch(to_keys([5, 6, 7]), _EMPTY_REQ_CTX)
# touch [3, 6, 7] (move to end of LRU order)
cpu_manager.touch(to_keys([3, 6, 7]), _EMPTY_REQ_CTX)
# prepare store [7, 9] -> evicts [8] (oldest following previous touch)
prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX)
@@ -320,7 +320,7 @@ def test_cpu_manager():
prepare_store_output,
ExpectedPrepareStoreOutput(
keys_to_store=[9],
store_block_ids=[1],
store_block_ids=[3],
evicted_keys=[8],
),
)
@@ -335,7 +335,7 @@ def test_cpu_manager():
verify_events(
cpu_manager.take_events(),
expected_stores=({3, 4, 5}, {6, 7, 8}),
expected_evictions=({2, 3, 4}, {8}),
expected_evictions=({4, 5, 2}, {8}),
)
@@ -295,6 +295,8 @@ class TestTieringOffloadingManager:
self.manager.prepare_store(blocks, _CTX)
self.manager.complete_store(blocks, _CTX, success=True)
self._simulate_on_schedule_end()
# for secondary tiers to drain jobs, so primary tier's blocks are evictable.
self._simulate_on_schedule_end()
self.secondary_tier1.touch = MagicMock(wraps=self.secondary_tier1.touch)
self.secondary_tier2.touch = MagicMock(wraps=self.secondary_tier2.touch)
@@ -303,7 +305,7 @@ class TestTieringOffloadingManager:
self.manager.touch(blocks, _CTX)
# Verify touch was called on primary tier (check LRU order)
primary_keys = list(self.primary_tier._policy.blocks.keys())
primary_keys = list(self.primary_tier._policy.evictable_blocks.keys())
assert primary_keys[-3:] == list(reversed(blocks))
# Verify touch was propagated to all secondary tiers
@@ -0,0 +1,286 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Self-describing KV cache events for the offloading connector.
The OffloadingManager identifies an offloaded chunk only by its OffloadKey,
so its raw events carry no token ids, parent hash, or block size.
:class:`OffloadingEventsTracker` snapshots each chunk's full ``BlockStored``
payload while the ``Request`` is alive and publishes stores as block-granular
payloads: a chunk event may carry multiple constituent per-block hashes, and
evictions fan out to the same hashes. Chunks overlapping a non-chunk-aligned
shared prefix re-announce the shared hashes once per chunk; consumers are
expected to deduplicate (reference-count) repeated store/remove announcements
of the same hash. Opt-in via
``kv_connector_extra_config["self_describing_kv_events"]``; inert unless
KV cache events are enabled. See the PR description for the full design.
"""
from collections.abc import Iterable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, NamedTuple
from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
from vllm.logger import init_logger
from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
from vllm.v1.kv_cache_interface import (
KVCacheGroupSpec,
get_kv_cache_spec_kind,
get_kv_cache_spec_sliding_window,
)
from vllm.v1.kv_offload.base import (
OffloadingEvent,
OffloadingKVEventsConfig,
OffloadKey,
get_offload_block_hash,
get_offload_group_idx,
)
from vllm.v1.request import Request
if TYPE_CHECKING:
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
GroupOffloadConfig,
)
logger = init_logger(__name__)
class OffloadingEventGroupSpec(NamedTuple):
kv_cache_spec_kind: str | None
kv_cache_spec_sliding_window: int | None
def get_offloading_event_group_spec(
kv_cache_group: KVCacheGroupSpec,
) -> OffloadingEventGroupSpec:
kv_cache_spec = kv_cache_group.kv_cache_spec
return OffloadingEventGroupSpec(
kv_cache_spec_kind=get_kv_cache_spec_kind(kv_cache_spec).value,
kv_cache_spec_sliding_window=get_kv_cache_spec_sliding_window(kv_cache_spec),
)
@dataclass(slots=True)
class _OffloadEventMetadata:
"""BlockStored payload snapshot for one OffloadKey, captured at store
time and kept until the matching eviction event. ``medium`` is forwarded
from the OffloadingEvent."""
# The chunk's constituent block hashes; the last one is the OffloadKey.
block_hashes: tuple[BlockHash, ...]
parent_block_hash: BlockHash | None
token_ids: tuple[int, ...]
block_size: int
lora_id: int | None
lora_name: str | None
# Deferred: needs the same incremental curr_mm_idx handling as GPU events.
extra_keys: tuple[tuple[Any, ...] | None, ...] | None
group_idx: int
kv_cache_spec: OffloadingEventGroupSpec
class OffloadingEventsTracker:
"""Tracks offloaded chunks' KV event payloads from store to eviction.
The scheduler calls :meth:`record_store` from ``_build_store_jobs``
while the ``Request`` is available, and routes the manager's raw
:class:`OffloadingEvent` stream through :meth:`take_events`. All state
is bounded by the CPU pool capacity and cleared by :meth:`reset`.
"""
def __init__(self, config: OffloadingKVEventsConfig):
self.config = config
self.self_describing_enabled = (
config.enable_kv_cache_events and config.self_describing_kv_events
)
# OffloadKey -> payload snapshot, kept until the eviction event so
# BlockRemoved can fan out. Bounded: one entry per offloaded chunk.
self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {}
def record_store(
self,
req: Request,
group_config: "GroupOffloadConfig",
offload_block_idx: int,
offload_key: OffloadKey,
) -> None:
"""Snapshot the KV cache event payload for one offloaded chunk.
No-op when self-describing event capture is disabled or for
sliding-window / SSM groups, which keep the legacy placeholder payload.
"""
if not self.self_describing_enabled:
return
if group_config.sliding_window_size_in_blocks is not None:
return
meta = self._build_event_metadata(req, group_config, offload_block_idx)
self._pending_event_metadata[offload_key] = meta
def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]:
"""Translate raw OffloadingEvents into self-describing KV events.
Complete metadata is available only for full-attention groups when
the tracker is enabled. Other shapes retain the legacy placeholder
payload so consumers can ignore them.
Yields:
``BlockStored`` or ``BlockRemoved`` events corresponding to
the underlying :class:`OffloadingEvent` stream.
"""
for event in events:
if event.removed:
yield from self._take_removed_event(event)
else:
yield from self._take_stored_event(event)
def reset(self) -> None:
"""Drop all tracked state; pending payloads are stale after a
manager cache reset."""
self._pending_event_metadata.clear()
def _build_event_metadata(
self,
req: Request,
group_config: "GroupOffloadConfig",
offload_block_idx: int,
) -> _OffloadEventMetadata:
"""Build the payload snapshot for one offloaded chunk: its
constituent per-block hashes, the whole chunk's tokens, and the
per-block ``block_size``."""
hbf = group_config.hash_block_size_factor
assert hbf > 0
assert offload_block_idx >= 0
# per-block token count (= the GPU/hash block size)
sub_block_size = group_config.offloaded_block_size // hbf
# chunk c covers hash-blocks [c*hbf, (c+1)*hbf); its tail block's hash
# is the chunk's OffloadKey.
first_hash_idx = offload_block_idx * hbf
last_hash_idx = first_hash_idx + hbf
assert first_hash_idx >= 0
assert last_hash_idx <= len(req.block_hashes)
chunk_hashes: list[BlockHash] = []
for block_hash in req.block_hashes[first_hash_idx:last_hash_idx]:
assert block_hash is not None
chunk_hashes.append(block_hash)
assert len(chunk_hashes) == hbf
if group_config.sliding_window_size_in_blocks is not None:
# record_store filters these out before calling this helper.
raise AssertionError("self-describing events only support full attention")
parent_block_hash: BlockHash | None
if first_hash_idx == 0:
parent_block_hash = None
else:
parent_block_hash = req.block_hashes[first_hash_idx - 1]
assert parent_block_hash is not None
tok_start = offload_block_idx * group_config.offloaded_block_size
tok_end = tok_start + group_config.offloaded_block_size
assert tok_end <= len(req.all_token_ids)
token_ids = tuple(req.all_token_ids[tok_start:tok_end])
lora_id: int | None = None
lora_name: str | None = None
if req.lora_request is not None:
lora_id = req.lora_request.adapter_id
lora_name = req.lora_request.name
return _OffloadEventMetadata(
block_hashes=tuple(chunk_hashes),
parent_block_hash=parent_block_hash,
token_ids=token_ids,
block_size=sub_block_size,
lora_id=lora_id,
lora_name=lora_name,
extra_keys=None,
group_idx=group_config.group_idx,
kv_cache_spec=group_config.kv_event_group_spec,
)
def _placeholder_stored(self, key: OffloadKey, medium: str) -> BlockStored:
return BlockStored(
block_hashes=[
maybe_convert_block_hash(BlockHash(get_offload_block_hash(key)))
],
parent_block_hash=None,
token_ids=[],
lora_id=None,
block_size=0,
medium=medium,
lora_name=None,
group_idx=get_offload_group_idx(key),
)
def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
# Metadata is read, NOT popped: the entry must survive until the
# eviction event so BlockRemoved can fan out to the same hashes.
# Events are self-contained (own parent), so key order is free.
for key in event.keys:
meta = self._pending_event_metadata.get(key)
if meta is None:
if self.self_describing_enabled:
# Expected for unsupported shapes; warn once only.
logger.warning_once(
"OffloadingEventsTracker: no event metadata for "
"offload key during BlockStored emission; emitting a "
"placeholder payload. Expected for non-full-attention "
"groups; otherwise indicates a missing populate path."
)
yield self._placeholder_stored(key, event.medium)
continue
yield BlockStored(
block_hashes=list(
maybe_convert_block_hash(h) for h in meta.block_hashes
),
parent_block_hash=(
maybe_convert_block_hash(meta.parent_block_hash)
if meta.parent_block_hash is not None
else None
),
token_ids=list(meta.token_ids),
block_size=meta.block_size,
lora_id=meta.lora_id,
medium=event.medium,
lora_name=meta.lora_name,
extra_keys=(
list(meta.extra_keys) if meta.extra_keys is not None else None
),
group_idx=meta.group_idx,
kv_cache_spec_kind=meta.kv_cache_spec.kv_cache_spec_kind,
kv_cache_spec_sliding_window=(
meta.kv_cache_spec.kv_cache_spec_sliding_window
),
)
def _take_removed_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
# Keep group_idx unambiguous if a manager batch spans groups.
by_group: dict[int, list] = {}
for key in event.keys:
meta = self._pending_event_metadata.pop(key, None)
if meta is not None:
group_idx = meta.group_idx
by_group.setdefault(group_idx, []).extend(
maybe_convert_block_hash(h) for h in meta.block_hashes
)
else:
if self.self_describing_enabled:
logger.warning_once(
"OffloadingEventsTracker: no event metadata for "
"offload key during BlockRemoved emission; emitting a "
"placeholder removal. Expected if the matching store "
"used the legacy placeholder payload; otherwise "
"indicates missing store metadata."
)
group_idx = get_offload_group_idx(key)
by_group.setdefault(group_idx, []).append(
maybe_convert_block_hash(BlockHash(get_offload_block_hash(key)))
)
for group_idx, hashes in by_group.items():
yield BlockRemoved(
block_hashes=hashes,
medium=event.medium,
group_idx=group_idx,
)
@@ -5,7 +5,7 @@ from dataclasses import dataclass, field
from itertools import islice
from typing import Any, NamedTuple
from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
from vllm.distributed.kv_events import KVCacheEvent
from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
@@ -14,6 +14,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
ReqId,
TransferJob,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import (
OffloadingEventGroupSpec,
OffloadingEventsTracker,
get_offloading_event_group_spec,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
OffloadingConnectorStats,
_TransferMetricName,
@@ -36,7 +41,6 @@ from vllm.v1.kv_offload.base import (
OffloadPolicy,
ReqContext,
RequestOffloadingContext,
get_offload_block_hash,
make_offload_key,
)
from vllm.v1.outputs import KVConnectorOutput
@@ -69,6 +73,9 @@ class GroupOffloadConfig(NamedTuple):
gpu_block_size: int
offloaded_block_size: int
hash_block_size_factor: int
# KV cache spec metadata propagated onto emitted BlockStored events so
# KV-aware consumers can classify and filter the group.
kv_event_group_spec: OffloadingEventGroupSpec
# None below means full attention
sliding_window_size_in_blocks: int | None
# Number of this group's offloaded blocks per full-attention alignment
@@ -200,6 +207,9 @@ class SchedulerOffloadConfig(NamedTuple):
alignment_block_count=_alignment_block_count(
gpu_block_size * spec.block_size_factor, sw
),
kv_event_group_spec=get_offloading_event_group_spec(
spec.kv_cache_config.kv_cache_groups[idx]
),
is_eagle_group=idx in eagle_groups,
)
for idx, gpu_block_size in enumerate(spec.gpu_block_size)
@@ -361,6 +371,8 @@ class OffloadingConnectorScheduler:
# be freed before a request finishes).
self._block_id_to_pending_jobs: dict[int, set[int]] = {}
self._events_tracker = OffloadingEventsTracker(spec.kv_events_config)
def _generate_job_id(self) -> int:
job_id = self._job_counter
self._job_counter += 1
@@ -934,6 +946,11 @@ class OffloadingConnectorScheduler:
continue
offloaded_block_idx = start_block_idx + idx
self._events_tracker.record_store(
req, group_config, offloaded_block_idx, offload_key
)
gpu_block_idx = offloaded_block_idx * block_size_factor
for i in range(block_size_factor):
block_id = block_ids[gpu_block_idx + i]
@@ -1184,25 +1201,17 @@ class OffloadingConnectorScheduler:
return False, None
def take_events(self) -> Iterable[KVCacheEvent]:
"""Take the KV cache events from the connector.
"""Drain pending KV cache events.
Returns:
A list of KV cache events.
Complete metadata is available only when self-describing KV events
are enabled, and only for full-attention groups. Other shapes retain
the previous placeholder payload so consumers can ignore them.
Yields:
``BlockStored`` or ``BlockRemoved`` events corresponding to
the underlying :class:`OffloadingEvent` stream.
"""
for event in self.manager.take_events():
block_hashes = [get_offload_block_hash(key) for key in event.keys]
if event.removed:
yield BlockRemoved(block_hashes=block_hashes, medium=event.medium)
else:
yield BlockStored(
block_hashes=block_hashes,
parent_block_hash=None,
token_ids=[],
lora_id=None,
block_size=0,
medium=event.medium,
lora_name=None,
)
yield from self._events_tracker.take_events(self.manager.take_events())
def reset_cache(self) -> None:
"""Reset the offloading manager cache, evicting all stored blocks."""
@@ -1238,6 +1247,10 @@ class OffloadingConnectorScheduler:
self._jobs.clear()
self._block_id_to_pending_jobs.clear()
# The manager pool is empty; pending event payloads and announced
# reference counts are stale.
self._events_tracker.reset()
# Note: _current_batch_jobs_to_flush is intentionally NOT cleared.
# The load flush IDs collected above must be delivered to workers.
if self._blocks_being_loaded is not None:
+2 -2
View File
@@ -61,7 +61,7 @@ def translate_error_response(response: ErrorResponse) -> JSONResponse:
async def create_messages(request: AnthropicMessagesRequest, raw_request: Request):
handler = messages(raw_request)
if handler is None:
base_server = raw_request.app.state.openai_serving_tokenization
base_server = raw_request.app.state.serving_tokenization
error = base_server.create_error_response(
NotImplementedError("The model does not support Messages API")
)
@@ -107,7 +107,7 @@ async def create_messages(request: AnthropicMessagesRequest, raw_request: Reques
async def count_tokens(request: AnthropicCountTokensRequest, raw_request: Request):
handler = messages(raw_request)
if handler is None:
base_server = raw_request.app.state.openai_serving_tokenization
base_server = raw_request.app.state.serving_tokenization
error = base_server.create_error_response(
NotImplementedError("The model does not support Messages API")
)
+11 -6
View File
@@ -34,7 +34,7 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware
from vllm.entrypoints.serve.render.serving import OpenAIServingRender
from vllm.entrypoints.serve.sagemaker.api_router import sagemaker_standards_bootstrap
from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
from vllm.entrypoints.serve.utils.api_utils import (
cli_env_setup,
log_non_default_args,
@@ -376,8 +376,7 @@ async def init_app_state(
log_error_stack=args.log_error_stack,
)
state.openai_serving_tokenization = OpenAIServingTokenization(
engine_client,
state.serving_tokenization = ServingTokenization(
state.openai_serving_models,
state.openai_serving_render,
request_logger=request_logger,
@@ -461,9 +460,15 @@ async def init_render_app_state(
)
state.openai_serving_models = model_registry
# Expose tokenization via the render handler (no engine required).
state.openai_serving_tokenization = state.openai_serving_render
state.serving_tokenization = ServingTokenization(
model_registry,
state.openai_serving_render,
request_logger=request_logger,
chat_template=resolved_chat_template,
chat_template_content_format=args.chat_template_content_format,
default_chat_template_kwargs=args.default_chat_template_kwargs,
trust_request_chat_template=args.trust_request_chat_template,
)
state.vllm_config = vllm_config
# Disable stats logging — there is no engine to poll.
+10 -265
View File
@@ -5,25 +5,19 @@ import time
from collections.abc import Awaitable, Mapping
from dataclasses import dataclass, field
from http import HTTPStatus
from typing import Any, ClassVar, Generic, Protocol, TypeAlias, TypeVar
from typing import ClassVar, Generic, TypeVar
from fastapi import Request
from pydantic import ConfigDict
from starlette.datastructures import Headers
import vllm.envs as envs
from vllm.config import ModelConfig
from vllm.engine.protocol import EngineClient
from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin
from vllm.entrypoints.openai.chat_completion.protocol import (
BatchChatCompletionRequest,
ChatCompletionRequest,
ChatCompletionResponse,
)
from vllm.entrypoints.openai.completion.protocol import (
CompletionRequest,
CompletionResponse,
)
from vllm.entrypoints.openai.engine.protocol import (
ErrorResponse,
@@ -31,81 +25,22 @@ from vllm.entrypoints.openai.engine.protocol import (
)
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.entrypoints.serve.disagg.protocol import GenerateRequest, GenerateResponse
from vllm.entrypoints.serve.tokenize.protocol import (
DetokenizeRequest,
TokenizeChatRequest,
TokenizeCompletionRequest,
TokenizeResponse,
)
from vllm.entrypoints.serve.utils.error_response import create_error_response
from vllm.entrypoints.serve.engine.serving import BaseServing
from vllm.entrypoints.serve.engine.typing import AnyRequest
from vllm.entrypoints.serve.utils.request_logger import RequestLogger
from vllm.entrypoints.speech_to_text.transcription.protocol import (
TranscriptionRequest,
TranscriptionResponse,
)
from vllm.entrypoints.speech_to_text.translation.protocol import TranslationRequest
from vllm.inputs import EngineInput, PromptType
from vllm.inputs import EngineInput
from vllm.logger import init_logger
from vllm.logprobs import Logprob, PromptLogprobs
from vllm.lora.request import LoRARequest
from vllm.renderers import ChatParams, TokenizeParams
from vllm.renderers.inputs.preprocess import (
extract_prompt_components,
extract_prompt_len,
)
from vllm.sampling_params import BeamSearchParams, SamplingParams
from vllm.tokenizers import TokenizerLike
from vllm.tracing import (
contains_trace_headers,
extract_trace_headers,
log_tracing_disabled_warning,
)
from vllm.utils import random_uuid
logger = init_logger(__name__)
class RendererRequest(Protocol):
def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
raise NotImplementedError
class RendererChatRequest(RendererRequest, Protocol):
def build_chat_params(
self,
default_template: str | None,
default_template_content_format: ChatTemplateContentFormatOption,
) -> ChatParams:
raise NotImplementedError
CompletionLikeRequest: TypeAlias = (
CompletionRequest | TokenizeCompletionRequest | DetokenizeRequest
)
ChatLikeRequest: TypeAlias = (
ChatCompletionRequest | BatchChatCompletionRequest | TokenizeChatRequest
)
SpeechToTextRequest: TypeAlias = TranscriptionRequest | TranslationRequest
AnyRequest: TypeAlias = (
CompletionLikeRequest
| ChatLikeRequest
| SpeechToTextRequest
| ResponsesRequest
| GenerateRequest
)
AnyResponse: TypeAlias = (
CompletionResponse
| ChatCompletionResponse
| TranscriptionResponse
| TokenizeResponse
| GenerateResponse
)
RequestT = TypeVar("RequestT", bound=AnyRequest)
_T = TypeVar("_T")
@@ -122,7 +57,7 @@ class ServeContext(Generic[RequestT]):
model_config = ConfigDict(arbitrary_types_allowed=True)
class OpenAIServing(BeamSearchOnlineMixin):
class OpenAIServing(BaseServing, BeamSearchOnlineMixin):
request_id_prefix: ClassVar[str] = """
A short string prepended to every requests ID.
"""
@@ -135,15 +70,14 @@ class OpenAIServing(BeamSearchOnlineMixin):
request_logger: RequestLogger | None,
return_tokens_as_token_ids: bool = False,
):
super().__init__()
super().__init__(
models=models,
model_config=engine_client.model_config,
request_logger=request_logger,
)
self.engine_client = engine_client
self.models = models
self.request_logger = request_logger
self.return_tokens_as_token_ids = return_tokens_as_token_ids
self.model_config = engine_client.model_config
self.renderer = engine_client.renderer
self.input_processor = engine_client.input_processor
vllm_config = getattr(engine_client, "vllm_config", None)
@@ -163,15 +97,6 @@ class OpenAIServing(BeamSearchOnlineMixin):
# Never fail server startup over the fingerprint.
self.system_fingerprint = None
@staticmethod
def create_error_response(
message: str | Exception,
err_type: str = "BadRequestError",
status_code: HTTPStatus = HTTPStatus.BAD_REQUEST,
param: str | None = None,
) -> ErrorResponse:
return create_error_response(message, err_type, status_code, param)
def create_streaming_error_response(
self,
message: str | Exception,
@@ -208,167 +133,6 @@ class OpenAIServing(BeamSearchOnlineMixin):
status_code=e.status_code,
)
async def _check_model(
self,
request: AnyRequest,
) -> ErrorResponse | None:
error_response = None
if self._is_model_supported(request.model):
return None
if request.model in self.models.lora_requests:
return None
if (
envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING
and request.model
and (load_result := await self.models.resolve_lora(request.model))
):
if isinstance(load_result, LoRARequest):
return None
if (
isinstance(load_result, ErrorResponse)
and load_result.error.code == HTTPStatus.BAD_REQUEST.value
):
error_response = load_result
return error_response or self.create_error_response(
message=f"The model `{request.model}` does not exist.",
err_type="NotFoundError",
status_code=HTTPStatus.NOT_FOUND,
param="model",
)
def _get_active_default_mm_loras(self, request: AnyRequest) -> LoRARequest | None:
"""Determine if there are any active default multimodal loras."""
# TODO: Currently this is only enabled for chat completions
# to be better aligned with only being enabled for .generate
# when run offline. It would be nice to support additional
# tasks types in the future.
message_types = self._get_message_types(request)
default_mm_loras = set()
for lora in self.models.lora_requests.values():
# Best effort match for default multimodal lora adapters;
# There is probably a better way to do this, but currently
# this matches against the set of 'types' in any content lists
# up until '_', e.g., to match audio_url -> audio
if lora.lora_name in message_types:
default_mm_loras.add(lora)
# Currently only support default modality specific loras if
# we have exactly one lora matched on the request.
if len(default_mm_loras) == 1:
return default_mm_loras.pop()
return None
def _maybe_get_adapters(
self,
request: AnyRequest,
supports_default_mm_loras: bool = False,
) -> LoRARequest | None:
if request.model in self.models.lora_requests:
return self.models.lora_requests[request.model]
# Currently only support default modality specific loras
# if we have exactly one lora matched on the request.
if supports_default_mm_loras:
default_mm_lora = self._get_active_default_mm_loras(request)
if default_mm_lora is not None:
return default_mm_lora
if self._is_model_supported(request.model):
return None
# if _check_model has been called earlier, this will be unreachable
raise ValueError(f"The model `{request.model}` does not exist.")
def _get_message_types(self, request: AnyRequest) -> set[str]:
"""Retrieve the set of types from message content dicts up
until `_`; we use this to match potential multimodal data
with default per modality loras.
"""
message_types: set[str] = set()
if not hasattr(request, "messages"):
return message_types
messages = request.messages
if messages is None or isinstance(messages, (str, bytes)):
return message_types
for message in messages:
if (
isinstance(message, dict)
and "content" in message
and isinstance(message["content"], list)
):
for content_dict in message["content"]:
if "type" in content_dict:
message_types.add(content_dict["type"].split("_")[0])
return message_types
def _validate_chat_template(
self,
request_chat_template: str | None,
chat_template_kwargs: dict[str, Any] | None,
trust_request_chat_template: bool,
) -> ErrorResponse | None:
if not trust_request_chat_template and (
request_chat_template is not None
or (
chat_template_kwargs
and chat_template_kwargs.get("chat_template") is not None
)
):
return self.create_error_response(
"Chat template is passed with request, but "
"--trust-request-chat-template is not set. "
"Refused request with untrusted chat template."
)
return None
@staticmethod
def _prepare_extra_chat_template_kwargs(
request_chat_template_kwargs: dict[str, Any] | None = None,
default_chat_template_kwargs: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Helper to merge server-default and request-specific chat template kwargs."""
request_chat_template_kwargs = request_chat_template_kwargs or {}
if default_chat_template_kwargs is None:
return request_chat_template_kwargs
# Apply server defaults first, then request kwargs override.
return default_chat_template_kwargs | request_chat_template_kwargs
def _extract_prompt_components(self, prompt: PromptType | EngineInput):
return extract_prompt_components(self.model_config, prompt)
def _extract_prompt_text(self, prompt: PromptType | EngineInput):
return self._extract_prompt_components(prompt).text
def _extract_prompt_len(self, prompt: EngineInput):
return extract_prompt_len(self.model_config, prompt)
def _log_inputs(
self,
request_id: str,
inputs: PromptType | EngineInput,
params: SamplingParams | BeamSearchParams | None,
lora_request: LoRARequest | None,
) -> None:
if self.request_logger is None:
return
components = self._extract_prompt_components(inputs)
self.request_logger.log_inputs(
request_id,
components.text,
components.token_ids,
components.embeds,
params=params,
lora_request=lora_request,
)
async def _get_trace_headers(
self,
headers: Headers,
@@ -383,18 +147,6 @@ class OpenAIServing(BeamSearchOnlineMixin):
return None
@staticmethod
def _base_request_id(
raw_request: Request | None, default: str | None = None
) -> str | None:
"""Pulls the request id to use from a header, if provided"""
if raw_request is not None and (
(req_id := raw_request.headers.get("X-Request-Id")) is not None
):
return req_id
return random_uuid() if default is None else default
@staticmethod
def _get_data_parallel_rank(raw_request: Request | None) -> int | None:
"""Pulls the data parallel rank from a header, if provided"""
@@ -464,13 +216,6 @@ class OpenAIServing(BeamSearchOnlineMixin):
return tokenizer.decode([token_id])
def _is_model_supported(self, model_name: str | None) -> bool:
if not model_name:
return True
if envs.VLLM_SKIP_MODEL_NAME_VALIDATION:
return True
return self.models.is_base_model(model_name)
def format_token_id_placeholder(token_id: int) -> str:
return f"token_id:{token_id}"
@@ -42,6 +42,10 @@ class OpenAIModelRegistry:
) -> None:
self.model_config = model_config
self.base_model_paths = base_model_paths
self.lora_requests: dict[str, LoRARequest] = {}
def model_name(self, lora_request: LoRARequest | None = None) -> str:
return self.base_model_paths[0].name
def is_base_model(self, model_name: str) -> bool:
return any(model.name == model_name for model in self.base_model_paths)
@@ -72,6 +76,9 @@ class OpenAIModelRegistry:
]
)
async def resolve_lora(self, lora_name: str):
raise RuntimeError("The OpenAIModelRegistry has no LoRA support.")
class OpenAIServingModels:
"""Shared instance to hold data about the loaded base model(s) and adapters.
@@ -12,7 +12,7 @@ from vllm.entrypoints.chat_utils import (
ChatTemplateContentFormatOption,
ConversationMessage,
)
from vllm.entrypoints.openai.engine.serving import RendererChatRequest, RendererRequest
from vllm.entrypoints.serve.engine.typing import RendererChatRequest, RendererRequest
from vllm.inputs import EngineInput, SingletonPrompt
from vllm.renderers import BaseRenderer, TokenizeParams, merge_kwargs
from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
+2 -2
View File
@@ -52,7 +52,7 @@ class PoolingServingBase(ABC):
self.engine_client = engine_client
self.models = models
self.model_config = models.model_config
self.renderer = models.renderer
self.renderer = engine_client.renderer
self.vllm_config = engine_client.vllm_config
self.max_model_len = self.model_config.max_model_len
self.request_logger = request_logger
@@ -61,7 +61,7 @@ class PoolingServingBase(ABC):
self.chat_template_config = chat_template_config
# Shared thread pool executor for preprocessing and postprocessing.
self._executor: Executor = models.renderer._executor
self._executor: Executor = self.renderer._executor
self._preprocessing_async = make_async(
self._preprocessing, executor=self._executor
)
+3 -3
View File
@@ -20,7 +20,7 @@ from vllm.entrypoints.serve.disagg.protocol import (
from vllm.entrypoints.serve.disagg.serving import (
ServingTokens,
)
from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
from vllm.entrypoints.serve.utils.api_utils import (
load_aware_call,
validate_json_request,
@@ -31,8 +31,8 @@ from vllm.logger import init_logger
logger = init_logger(__name__)
def tokenization(request: Request) -> OpenAIServingTokenization:
return request.app.state.openai_serving_tokenization
def tokenization(request: Request) -> ServingTokenization:
return request.app.state.serving_tokenization
def generate_tokens(request: Request) -> ServingTokens | None:
+193
View File
@@ -0,0 +1,193 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from http import HTTPStatus
from fastapi import Request
from vllm import PromptType, SamplingParams, envs
from vllm.config import ModelConfig
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
from vllm.entrypoints.openai.models.serving import (
OpenAIModelRegistry,
OpenAIServingModels,
)
from vllm.entrypoints.serve.engine.typing import AnyRequest
from vllm.entrypoints.serve.utils.error_response import create_error_response
from vllm.entrypoints.serve.utils.request_logger import RequestLogger
from vllm.inputs import EngineInput
from vllm.lora.request import LoRARequest
from vllm.renderers.inputs.preprocess import (
extract_prompt_components,
extract_prompt_len,
)
from vllm.sampling_params import BeamSearchParams
from vllm.utils import random_uuid
class BaseServing:
def __init__(
self,
models: OpenAIServingModels | OpenAIModelRegistry,
model_config: ModelConfig,
request_logger: RequestLogger | None = None,
):
self.models = models
self.model_config = model_config
self.request_logger = request_logger
async def _check_model(
self,
request: AnyRequest,
) -> ErrorResponse | None:
error_response = None
if self._is_model_supported(request.model):
return None
if request.model in self.models.lora_requests:
return None
if (
envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING
and request.model
and (load_result := await self.models.resolve_lora(request.model))
):
if isinstance(load_result, LoRARequest):
return None
if (
isinstance(load_result, ErrorResponse)
and load_result.error.code == HTTPStatus.BAD_REQUEST.value
):
error_response = load_result
return error_response or self.create_error_response(
message=f"The model `{request.model}` does not exist.",
err_type="NotFoundError",
status_code=HTTPStatus.NOT_FOUND,
param="model",
)
def _is_model_supported(self, model_name: str | None) -> bool:
if not model_name:
return True
if envs.VLLM_SKIP_MODEL_NAME_VALIDATION:
return True
return self.models.is_base_model(model_name)
@staticmethod
def create_error_response(
message: str | Exception,
err_type: str = "BadRequestError",
status_code: HTTPStatus = HTTPStatus.BAD_REQUEST,
param: str | None = None,
) -> ErrorResponse:
return create_error_response(message, err_type, status_code, param)
def _extract_prompt_components(self, prompt: PromptType | EngineInput):
return extract_prompt_components(self.model_config, prompt)
def _extract_prompt_text(self, prompt: PromptType | EngineInput):
return self._extract_prompt_components(prompt).text
def _extract_prompt_len(self, prompt: EngineInput):
return extract_prompt_len(self.model_config, prompt)
def _log_inputs(
self,
request_id: str,
inputs: PromptType | EngineInput,
params: SamplingParams | BeamSearchParams | None,
lora_request: LoRARequest | None,
) -> None:
if self.request_logger is None:
return
components = self._extract_prompt_components(inputs)
self.request_logger.log_inputs(
request_id,
components.text,
components.token_ids,
components.embeds,
params=params,
lora_request=lora_request,
)
@staticmethod
def _base_request_id(
raw_request: Request | None, default: str | None = None
) -> str | None:
"""Pulls the request id to use from a header, if provided"""
if raw_request is not None and (
(req_id := raw_request.headers.get("X-Request-Id")) is not None
):
return req_id
return random_uuid() if default is None else default
def _get_message_types(self, request: AnyRequest) -> set[str]:
"""Retrieve the set of types from message content dicts up
until `_`; we use this to match potential multimodal data
with default per modality loras.
"""
message_types: set[str] = set()
if not hasattr(request, "messages"):
return message_types
messages = request.messages
if messages is None or isinstance(messages, (str, bytes)):
return message_types
for message in messages:
if (
isinstance(message, dict)
and "content" in message
and isinstance(message["content"], list)
):
for content_dict in message["content"]:
if "type" in content_dict:
message_types.add(content_dict["type"].split("_")[0])
return message_types
def _get_active_default_mm_loras(self, request: AnyRequest) -> LoRARequest | None:
"""Determine if there are any active default multimodal loras."""
# TODO: Currently this is only enabled for chat completions
# to be better aligned with only being enabled for .generate
# when run offline. It would be nice to support additional
# tasks types in the future.
message_types = self._get_message_types(request)
default_mm_loras = set()
for lora in self.models.lora_requests.values():
# Best effort match for default multimodal lora adapters;
# There is probably a better way to do this, but currently
# this matches against the set of 'types' in any content lists
# up until '_', e.g., to match audio_url -> audio
if lora.lora_name in message_types:
default_mm_loras.add(lora)
# Currently only support default modality specific loras if
# we have exactly one lora matched on the request.
if len(default_mm_loras) == 1:
return default_mm_loras.pop()
return None
def _maybe_get_adapters(
self,
request: AnyRequest,
supports_default_mm_loras: bool = False,
) -> LoRARequest | None:
if request.model in self.models.lora_requests:
return self.models.lora_requests[request.model]
# Currently only support default modality specific loras
# if we have exactly one lora matched on the request.
if supports_default_mm_loras:
default_mm_lora = self._get_active_default_mm_loras(request)
if default_mm_lora is not None:
return default_mm_lora
if self._is_model_supported(request.model):
return None
# if _check_model has been called earlier, this will be unreachable
raise ValueError(f"The model `{request.model}` does not exist.")
+71
View File
@@ -0,0 +1,71 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Protocol, TypeAlias
from vllm.config import ModelConfig
from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
from vllm.entrypoints.openai.chat_completion.protocol import (
BatchChatCompletionRequest,
ChatCompletionRequest,
ChatCompletionResponse,
)
from vllm.entrypoints.openai.completion.protocol import (
CompletionRequest,
CompletionResponse,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.entrypoints.serve.disagg.protocol import GenerateRequest, GenerateResponse
from vllm.entrypoints.serve.tokenize.protocol import (
DetokenizeRequest,
TokenizeChatRequest,
TokenizeCompletionRequest,
TokenizeResponse,
)
from vllm.entrypoints.speech_to_text.transcription.protocol import (
TranscriptionRequest,
TranscriptionResponse,
)
from vllm.entrypoints.speech_to_text.translation.protocol import TranslationRequest
from vllm.renderers import ChatParams, TokenizeParams
class RendererRequest(Protocol):
def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
raise NotImplementedError
class RendererChatRequest(RendererRequest, Protocol):
def build_chat_params(
self,
default_template: str | None,
default_template_content_format: ChatTemplateContentFormatOption,
) -> ChatParams:
raise NotImplementedError
CompletionLikeRequest: TypeAlias = (
CompletionRequest | TokenizeCompletionRequest | DetokenizeRequest
)
ChatLikeRequest: TypeAlias = (
ChatCompletionRequest | BatchChatCompletionRequest | TokenizeChatRequest
)
SpeechToTextRequest: TypeAlias = TranscriptionRequest | TranslationRequest
AnyRequest: TypeAlias = (
CompletionLikeRequest
| ChatLikeRequest
| SpeechToTextRequest
| ResponsesRequest
| GenerateRequest
)
AnyResponse: TypeAlias = (
CompletionResponse
| ChatCompletionResponse
| TranscriptionResponse
| TokenizeResponse
| GenerateResponse
)
@@ -5,8 +5,7 @@ from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse
from vllm.engine.protocol import EngineClient
from vllm.entrypoints.openai.engine.serving import OpenAIServing
from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
from vllm.logger import init_logger
from vllm.version import __version__ as VLLM_VERSION
@@ -15,13 +14,13 @@ router = APIRouter()
logger = init_logger(__name__)
def base(request: Request) -> OpenAIServing:
def base(request: Request) -> ServingTokenization:
# Reuse the existing instance
return tokenization(request)
def tokenization(request: Request) -> OpenAIServingTokenization:
return request.app.state.openai_serving_tokenization
def tokenization(request: Request) -> ServingTokenization:
return request.app.state.serving_tokenization
def engine_client(request: Request) -> EngineClient:
@@ -9,16 +9,14 @@ from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from typing_extensions import assert_never
from vllm.entrypoints.openai.engine.protocol import (
ErrorResponse,
)
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
from vllm.entrypoints.serve.tokenize.protocol import (
DetokenizeRequest,
DetokenizeResponse,
TokenizeRequest,
TokenizeResponse,
)
from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
from vllm.entrypoints.serve.utils.api_utils import (
validate_json_request,
with_cancellation,
@@ -28,8 +26,8 @@ from vllm.logger import init_logger
logger = init_logger(__name__)
def tokenization(request: Request) -> OpenAIServingTokenization:
return request.app.state.openai_serving_tokenization
def tokenization(request: Request) -> ServingTokenization:
return request.app.state.serving_tokenization
router = APIRouter()
+9 -7
View File
@@ -5,11 +5,13 @@ from typing import Any, Final
from fastapi import Request
from vllm.engine.protocol import EngineClient
from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
from vllm.entrypoints.openai.engine.serving import OpenAIServing
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.openai.models.serving import (
OpenAIModelRegistry,
OpenAIServingModels,
)
from vllm.entrypoints.serve.engine.serving import BaseServing
from vllm.entrypoints.serve.render.serving import OpenAIServingRender
from vllm.entrypoints.serve.tokenize.protocol import (
DetokenizeRequest,
@@ -27,11 +29,10 @@ from vllm.tokenizers import TokenizerLike
logger = init_logger(__name__)
class OpenAIServingTokenization(OpenAIServing):
class ServingTokenization(BaseServing):
def __init__(
self,
engine_client: EngineClient,
models: OpenAIServingModels,
models: OpenAIServingModels | OpenAIModelRegistry,
openai_serving_render: OpenAIServingRender,
*,
request_logger: RequestLogger | None,
@@ -41,11 +42,12 @@ class OpenAIServingTokenization(OpenAIServing):
trust_request_chat_template: bool = False,
) -> None:
super().__init__(
engine_client=engine_client,
models=models,
model_config=openai_serving_render.model_config,
request_logger=request_logger,
)
self.renderer = openai_serving_render.renderer
self.openai_serving_render = openai_serving_render
self.chat_template = chat_template
self.chat_template_content_format: Final = chat_template_content_format
@@ -22,8 +22,9 @@ from vllm.entrypoints.openai.engine.protocol import (
RequestResponseMetadata,
UsageInfo,
)
from vllm.entrypoints.openai.engine.serving import OpenAIServing, SpeechToTextRequest
from vllm.entrypoints.openai.engine.serving import OpenAIServing
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.serve.engine.typing import SpeechToTextRequest
from vllm.entrypoints.serve.utils.api_utils import get_max_tokens
from vllm.entrypoints.serve.utils.request_logger import RequestLogger
from vllm.exceptions import VLLMValidationError
@@ -581,7 +581,14 @@ class Attention(nn.Module, AttentionLayerBase):
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
# Block size may get updated after model loading, refresh it
block_size = vllm_config.cache_config.block_size
# Should not be called for enc-dec or encoder-only attention.
# Encoder-only attention is prefill-only and keeps no autoregressive KV
# cache. In hybrid models (e.g. Qwen3.5 / ColQwen3.5: GatedDeltaNet
# linear_attention interleaved with full_attention) the runner iterates
# every attention module to build the KV-cache spec, so an ENCODER_ONLY
# full_attention layer reaches here; it contributes no KV cache group.
if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
return None
# Should not be called for enc-dec attention.
assert self.attn_type == AttentionType.DECODER
quant_mode = get_kv_quant_mode(self.kv_cache_dtype)
if self.sliding_window is not None:
@@ -67,6 +67,7 @@ class XPUExperts(mk.FusedMoEExpertsModular):
self.is_mxfp4 = False
self.is_block_fp8 = False
self.is_mxfp8 = False
self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
self.fused_moe_impl: XpuFusedMoe | None = None
@property
@@ -176,6 +177,7 @@ class XPUExperts(mk.FusedMoEExpertsModular):
is_mxfp4=self.is_mxfp4,
is_mxfp8=self.is_mxfp8,
is_block_fp8=self.is_block_fp8,
gemm1_clamp_limit=self.gemm1_clamp_limit,
)
assert self.fused_moe_impl is not None
self.fused_moe_impl.apply(
+40 -5
View File
@@ -6,9 +6,25 @@ import torch
# import vllm.model_executor.kernels.mhc # noqa: F401
import vllm.model_executor.kernels.mhc as mhc_kernels
from vllm.model_executor.custom_op import CustomOp
from vllm.platforms import current_platform
from vllm.utils.import_utils import has_tilelang
HAS_TILELANG = has_tilelang()
def _has_tilelang_mhc() -> bool:
if not has_tilelang():
return False
if current_platform.is_cuda():
return True
if current_platform.is_rocm():
from vllm.platforms.rocm import on_gfx942
# TileLang MHC currently produces incorrect results on gfx942. Keep
# gfx942 on the existing torch/triton fallbacks until that path is fixed.
return not on_gfx942()
return False
HAS_TILELANG_MHC = _has_tilelang_mhc()
# --8<-- [start:mhc_pre]
@@ -89,7 +105,7 @@ class MHCPreOp(CustomOp):
# sinkhorn_repeat,
# )
# else:
if HAS_TILELANG:
if HAS_TILELANG_MHC:
return torch.ops.vllm.mhc_pre_tilelang(
residual,
fn,
@@ -224,7 +240,7 @@ class MHCPostOp(CustomOp):
# comb_res_mix,
# )
# else:
if HAS_TILELANG:
if HAS_TILELANG_MHC:
return torch.ops.vllm.mhc_post_tilelang(
x, residual, post_layer_mix, comb_res_mix
)
@@ -310,7 +326,7 @@ class HCHeadOp(CustomOp):
outer_shape = hidden_states.shape[:-2]
hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
if HAS_TILELANG:
if HAS_TILELANG_MHC:
out = torch.ops.vllm.hc_head_fused_kernel_tilelang(
hs_flat,
hc_fn,
@@ -447,7 +463,26 @@ class MHCFusedPostPreOp(CustomOp):
norm_weight: torch.Tensor | None = None,
norm_eps: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
return torch.ops.vllm.mhc_fused_post_pre_tilelang(
if HAS_TILELANG_MHC:
return torch.ops.vllm.mhc_fused_post_pre_tilelang(
x,
residual,
post_layer_mix,
comb_res_mix,
fn,
hc_scale,
hc_base,
rms_eps,
hc_pre_eps,
hc_sinkhorn_eps,
hc_post_mult_value,
sinkhorn_repeat,
n_splits,
tile_n,
norm_weight,
norm_eps,
)
return self.forward_native(
x,
residual,
post_layer_mix,
@@ -426,6 +426,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
device = get_tp_group().device
tp_rank = get_tensor_model_parallel_rank()
tp_size = layer.moe_config.moe_parallel_config.tp_size
loaded_weight = loaded_weight.to(device)
shard_size = layer.intermediate_size_per_partition
@@ -464,9 +465,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
)
if "w13_qzeros" in weight_name:
tensor = loaded_weight.view(layer.tp_size, -1, loaded_weight.size(1))[
tp_rank
]
tensor = loaded_weight.view(tp_size, -1, loaded_weight.size(1))[tp_rank]
if shard_id == "w1":
param.data[expert_id, : shard_size // 2] = tensor
else:
@@ -474,7 +473,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
return True if return_success else None
elif "w2_qzeros" in weight_name:
param.data[expert_id] = loaded_weight.view(
loaded_weight.size(0), layer.tp_size, -1
loaded_weight.size(0), tp_size, -1
)[:, tp_rank]
return True if return_success else None
else:
+9 -1
View File
@@ -15,6 +15,7 @@ Based on: Qwen3.5 backbone with custom text projection
Target models:
- athrael-soju/colqwen3.5-4.5B-v3
- vultr/VultronRetrieverPrime-Qwen3.5-8B
"""
from collections.abc import Iterable, Mapping
@@ -166,12 +167,19 @@ class ColQwen3_5Model(
or 128 # default from reference implementation
)
# ColPali defines `custom_text_proj = nn.Linear(hidden, dim)`, i.e.
# bias=True by default, and the trained ColQwen3.5 checkpoints ship a
# `custom_text_proj.bias`. Construct with a bias and zero-initialize it:
# a (legacy) bias-less checkpoint then behaves identically to bias=False,
# while load_weights() below picks up a trained bias instead of silently
# dropping it (which shifts every per-token vector and the MaxSim ranking).
self.custom_text_proj = nn.Linear(
hidden_size,
self.embed_dim,
bias=False,
bias=True,
dtype=head_dtype,
)
nn.init.zeros_(self.custom_text_proj.bias)
pooler_config = vllm_config.model_config.pooler_config
assert pooler_config is not None
+15 -1
View File
@@ -627,6 +627,20 @@ class Qwen3_5ForConditionalGenerationConfig(VerifyAndUpdateConfig):
)
class ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig):
"""ColQwen3.5 (late-interaction retrieval) inherits Qwen3.5's mamba cache
handling and additionally serves BIDIRECTIONAL attention: ColPali-style
document/query encoding attends over the whole sequence, not causally. Set
is_causal=False so Qwen3NextAttention builds its full_attention layers with
AttentionType.ENCODER_ONLY (the linear_attention GatedDeltaNet layers are
unaffected). Generation arches keep the parent (causal) and are untouched.
"""
@staticmethod
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
model_config.hf_config.is_causal = False
class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):
@staticmethod
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -656,7 +670,7 @@ class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig):
MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
"ColBERTJinaRobertaModel": JinaRobertaModelConfig,
"ColQwen3_5": Qwen3_5ForConditionalGenerationConfig,
"ColQwen3_5": ColQwen3_5Config,
"DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig,
"DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
"DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501
File diff suppressed because it is too large Load Diff
-385
View File
@@ -1,385 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable, Mapping
from typing import Annotated, Literal, TypeAlias
import torch
import torch.nn as nn
from transformers import BatchFeature, PretrainedConfig
from transformers.models.llava_next.modeling_llava_next import (
get_anyres_image_grid_shape,
unpad_image,
)
from vllm.config import VllmConfig
from vllm.model_executor.layers.activation import get_act_fn
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import MultiModalFieldConfig
from vllm.sequence import IntermediateTensors
from vllm.utils.tensor_schema import TensorSchema, TensorShape
from .clip import CLIPVisionModel
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
from .llava import (
BaseLlavaMultiModalProcessor,
LlavaDummyInputsBuilder,
init_vision_tower_for_llava,
)
from .llava_next import LlavaNextProcessingInfo
from .pixtral import PixtralHFVisionModel
from .siglip import SiglipVisionModel
from .utils import (
AutoWeightsLoader,
init_vllm_registered_model,
maybe_prefix,
)
class MiniMaxVL01ImagePixelInputs(TensorSchema):
"""
Dimensions:
- bn: Batch size * number of images
- np: Number of patches + 1
- c: Number of channels (3)
- h: Height
- w: Width
Note that `num_patches` may be different per batch and image,
in which case the data is passed as a list instead of a batched tensor.
"""
type: Literal["pixel_values"] = "pixel_values"
pixel_values: Annotated[
torch.Tensor | list[torch.Tensor],
TensorShape("bn", "np", 3, "h", "w", dynamic_dims={"np", "h", "w"}),
]
image_sizes: Annotated[torch.Tensor | None, TensorShape("bn", 2)]
# This should be in `(height, width)` format.
class MiniMaxVL01ImageEmbeddingInputs(TensorSchema):
"""
Dimensions:
- bn: Batch size * number of images
- ifs: Image feature size
- hs: Hidden size (must match language model backbone)
"""
type: Literal["image_embeds"] = "image_embeds"
data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")]
MiniMaxVL01ImageInputs: TypeAlias = (
MiniMaxVL01ImagePixelInputs | MiniMaxVL01ImageEmbeddingInputs
)
class MiniMaxVL01MultiModalProjector(nn.Module):
def __init__(
self,
vision_hidden_size: int,
text_hidden_size: int,
projector_hidden_act: str,
multimodal_projector_bias: bool,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
):
super().__init__()
self.linear_1 = ColumnParallelLinear(
vision_hidden_size,
text_hidden_size,
bias=multimodal_projector_bias,
quant_config=quant_config,
prefix=f"{prefix}.linear_1",
)
self.act = get_act_fn(projector_hidden_act)
self.linear_2 = RowParallelLinear(
text_hidden_size,
text_hidden_size,
bias=multimodal_projector_bias,
quant_config=quant_config,
prefix=f"{prefix}.linear_2",
)
def forward(self, image_features: torch.Tensor) -> torch.Tensor:
hidden_states, _ = self.linear_1(image_features)
hidden_states = self.act(hidden_states)
hidden_states, _ = self.linear_2(hidden_states)
return hidden_states
class MiniMaxVL01DummyInputsBuilder(LlavaDummyInputsBuilder):
pass
class MiniMaxVL01ProcessingInfo(LlavaNextProcessingInfo):
def get_hf_config(self): # Need to override the config type
return self.ctx.get_hf_config(PretrainedConfig)
def get_hf_processor(self, **kwargs: object):
hf_processor = self.ctx.get_hf_processor(**kwargs)
image_processor = hf_processor.image_processor
image_processor.anyres_preprocess = image_processor.anyres_for_vllm_preprocess
return hf_processor
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
return {"image": None}
class MiniMaxVL01MultiModalProcessor(
BaseLlavaMultiModalProcessor[MiniMaxVL01ProcessingInfo]
):
def _call_hf_processor(
self,
prompt: str,
mm_data: Mapping[str, object],
mm_kwargs: Mapping[str, object],
tok_kwargs: Mapping[str, object],
) -> BatchFeature:
processed_outputs = super()._call_hf_processor(
prompt=prompt,
mm_data=mm_data,
mm_kwargs=mm_kwargs,
tok_kwargs=tok_kwargs,
)
pixel_values = processed_outputs.get("pixel_values")
if pixel_values is not None:
# Avoid padding since we need the output for each image to be
# independent of other images for the cache to work correctly
image_sizes = processed_outputs["image_sizes"]
assert len(pixel_values) == len(image_sizes)
processed_outputs["pixel_values"] = [
p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes)
]
return processed_outputs
def _get_mm_fields_config(
self,
hf_inputs: BatchFeature,
hf_processor_mm_kwargs: Mapping[str, object],
) -> Mapping[str, MultiModalFieldConfig]:
return {
"pixel_values": MultiModalFieldConfig.batched("image"),
"image_sizes": MultiModalFieldConfig.batched("image"),
"image_embeds": MultiModalFieldConfig.batched("image"),
}
@MULTIMODAL_REGISTRY.register_processor(
MiniMaxVL01MultiModalProcessor,
info=MiniMaxVL01ProcessingInfo,
dummy_inputs=MiniMaxVL01DummyInputsBuilder,
)
class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"gate_up_proj": ["gate_proj", "up_proj"],
}
@classmethod
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
if modality.startswith("image"):
return "<image>"
raise ValueError("Only image modality is supported")
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
super().__init__()
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
multimodal_config = vllm_config.model_config.multimodal_config
self.config = config
self.multimodal_config = multimodal_config
with self._mark_tower_model(vllm_config, "image"):
self.vision_tower = init_vision_tower_for_llava(
config,
quant_config=quant_config,
require_post_norm=False,
prefix=maybe_prefix(prefix, "vision_tower"),
)
self.multi_modal_projector = MiniMaxVL01MultiModalProjector(
vision_hidden_size=config.vision_config.hidden_size,
text_hidden_size=config.text_config.hidden_size,
projector_hidden_act=config.projector_hidden_act,
multimodal_projector_bias=True,
quant_config=quant_config,
prefix=maybe_prefix(prefix, "multi_modal_projector"),
)
self.image_newline = nn.Parameter(
torch.empty(config.text_config.hidden_size)
)
with self._mark_language_model(vllm_config):
self.language_model = init_vllm_registered_model(
vllm_config=vllm_config,
hf_config=config.text_config,
prefix=maybe_prefix(prefix, "language_model"),
)
self.vision_feature_layer = config.vision_feature_layer
self.vocab_size = config.text_config.vocab_size
self.pad_token_id = -1
if self.config.text_config.pad_token_id is not None:
self.pad_token_id = self.config.text_config.pad_token_id
self.make_empty_intermediate_tensors = (
self.language_model.make_empty_intermediate_tensors
)
def _image_pixels_to_features(
self,
vision_tower: CLIPVisionModel | SiglipVisionModel | PixtralHFVisionModel,
pixel_values: torch.Tensor | list[torch.Tensor],
) -> torch.Tensor | tuple[torch.Tensor, ...]:
# NOTE: we skip the step to select the vision feature layer since
# this is already done inside the vision tower
feature_select_strategy = self.config.vision_feature_select_strategy
return tuple(
vision_tower(p, feature_select_strategy=feature_select_strategy)
for p in pixel_values
)
# adapted from https://huggingface.co/MiniMaxAI/MiniMax-VL-01/blob/main/modeling_minimax_vl_01.py#L616-L631
def pack_image_features(
self, image_features: list[torch.Tensor], image_sizes: torch.Tensor
):
new_image_features = []
for image_idx, image_feature in enumerate(image_features):
if image_feature.shape[0] > 1:
base_image_feature = image_feature[0]
image_feature = image_feature[1:]
height = width = (
self.config.vision_config.image_size
// self.config.vision_config.patch_size
)
if height * width != base_image_feature.shape[0]:
raise ValueError(
"The number of patches is not consistent with the image size."
)
num_patch_height, num_patch_width = get_anyres_image_grid_shape(
image_sizes[image_idx],
self.config.image_grid_pinpoints,
self.config.vision_config.image_size,
)
image_feature = image_feature.view(
num_patch_height, num_patch_width, height, width, -1
)
image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
image_feature = image_feature.flatten(1, 2).flatten(2, 3)
image_feature = unpad_image(image_feature, image_sizes[image_idx])
image_feature = torch.cat(
(
image_feature,
self.image_newline[:, None, None]
.expand(*image_feature.shape[:-1], 1)
.to(image_feature.dtype),
),
dim=-1,
)
image_feature = image_feature.flatten(1, 2).transpose(0, 1)
image_feature = torch.cat((base_image_feature, image_feature), dim=0)
else:
image_feature = image_feature[0]
image_feature = torch.cat(
(image_feature, self.image_newline[None].to(image_feature)), dim=0
)
new_image_features.append(image_feature)
return new_image_features
def _process_image_pixels(
self,
inputs: MiniMaxVL01ImagePixelInputs,
) -> torch.Tensor | tuple[torch.Tensor, ...]:
pixel_values = inputs["pixel_values"]
return self._image_pixels_to_features(self.vision_tower, pixel_values)
def _process_image_input(
self,
image_input: MiniMaxVL01ImageInputs,
) -> torch.Tensor | tuple[torch.Tensor, ...]:
if image_input["type"] == "image_embeds":
return image_input["data"]
image_features = self._process_image_pixels(image_input)
if isinstance(image_features, torch.Tensor):
return self.multi_modal_projector(image_features)
feature_sizes = [image_feature.shape[0] for image_feature in image_features]
image_embeds = self.multi_modal_projector(torch.cat(image_features))
image_embeds = torch.split(image_embeds, feature_sizes)
image_sizes = image_input.get("image_sizes")
return self.pack_image_features(image_embeds, image_sizes)
def _parse_and_validate_image_input(
self, **kwargs: object
) -> MiniMaxVL01ImageInputs | None:
pixel_values = kwargs.pop("pixel_values", None)
image_sizes = kwargs.pop("image_sizes", None)
image_embeds = kwargs.pop("image_embeds", None)
if pixel_values is None and image_embeds is None:
return None
if pixel_values is not None and image_sizes is not None:
return MiniMaxVL01ImagePixelInputs(
type="pixel_values",
pixel_values=pixel_values,
image_sizes=image_sizes,
)
if image_embeds is not None:
return MiniMaxVL01ImageEmbeddingInputs(
type="image_embeds",
data=image_embeds,
)
raise AssertionError("This line should be unreachable.")
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
image_input = self._parse_and_validate_image_input(**kwargs)
if image_input is None:
return []
return self._process_image_input(image_input)
def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
**kwargs: object,
) -> torch.Tensor | IntermediateTensors:
if intermediate_tensors is not None:
inputs_embeds = None
hidden_states = self.language_model.model(
input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
)
return hidden_states
def compute_logits(
self,
hidden_states: torch.Tensor,
) -> torch.Tensor | None:
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self)
return loader.load_weights(weights)
+11
View File
@@ -62,6 +62,7 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
from vllm.v1.attention.backend import AttentionType
from .interfaces import (
EagleModelMixin,
@@ -267,6 +268,15 @@ class Qwen3NextAttention(nn.Module):
dual_chunk_attention_config=self.dual_chunk_attention_config,
)
# Late-interaction retrieval models (e.g. ColQwen3.5) run BIDIRECTIONAL
# attention on the full_attention layers; they set config.is_causal=False
# via a VerifyAndUpdateConfig handler. Generation models leave is_causal
# unset (-> causal/DECODER), so this is a no-op for them. Mirrors qwen3.py.
attn_type = (
AttentionType.DECODER
if getattr(config, "is_causal", True)
else AttentionType.ENCODER_ONLY
)
self.attn = Attention(
self.num_heads,
self.head_dim,
@@ -275,6 +285,7 @@ class Qwen3NextAttention(nn.Module):
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.attn",
attn_type=attn_type,
**{
"layer_idx": extract_layer_index(prefix),
"dual_chunk_attention_config": self.dual_chunk_attention_config,
+4 -7
View File
@@ -159,9 +159,6 @@ _TEXT_GENERATION_MODELS = {
"MellumForCausalLM": ("mellum", "MellumForCausalLM"),
"MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
"MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
"MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
"MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
"MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
"MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
"MiniMaxM3SparseForCausalLM": (
"vllm.models.minimax_m3",
@@ -490,10 +487,6 @@ _MULTIMODAL_MODELS = {
"vllm.models.minimax_m3",
"MiniMaxM3SparseForConditionalGeneration",
),
"MiniMaxVL01ForConditionalGeneration": (
"minimax_vl_01",
"MiniMaxVL01ForConditionalGeneration",
),
"MiniCPMO": ("minicpmo", "MiniCPMO"),
"MiniCPMV": ("minicpmv", "MiniCPMV"),
"MiniCPMV4_6ForConditionalGeneration": (
@@ -735,6 +728,10 @@ _PREVIOUSLY_SUPPORTED_MODELS = {
"XverseForCausalLM": "0.23.0",
"Dots1ForCausalLM": "0.23.0",
"BambaForCausalLM": "0.23.0",
"MiniMaxForCausalLM": "0.23.0",
"MiniMaxText01ForCausalLM": "0.23.0",
"MiniMaxM1ForCausalLM": "0.23.0",
"MiniMaxVL01ForConditionalGeneration": "0.23.0",
}
_OOT_SUPPORTED_MODELS = {
+3 -3
View File
@@ -27,6 +27,7 @@ from vllm.model_executor.layers.linear import (
)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.mhc import (
HAS_TILELANG_MHC,
HCHeadOp,
MHCFusedPostPreOp,
MHCPostOp,
@@ -51,7 +52,6 @@ from vllm.model_executor.models.utils import (
from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.utils.import_utils import has_tilelang
class DeepseekV4MLP(nn.Module):
@@ -303,7 +303,7 @@ class DeepseekV4DecoderLayer(nn.Module):
self.mhc_pre = MHCPreOp()
self.mhc_post = MHCPostOp()
self.mhc_fused_post_pre = MHCFusedPostPreOp()
self.has_tilelang = has_tilelang()
self.has_tilelang = HAS_TILELANG_MHC
def hc_pre(
self,
@@ -513,7 +513,7 @@ class DeepseekV4Model(nn.Module):
requires_grad=False,
)
self.hc_head_op = HCHeadOp()
self.has_tilelang = has_tilelang()
self.has_tilelang = HAS_TILELANG_MHC
# Pre-hc_head residual stream buffer for the MTP draft. Stable
# address (outside the cudagraph pool) so the copy_ in forward()
# refreshes it correctly across captured shapes.

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