forked from Karylab-cklius/vllm
Compare commits
19
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
cf30ef60cf | ||
|
|
89accad2cc | ||
|
|
3c8e49596c | ||
|
|
cec2ec1176 | ||
|
|
435f82d61a | ||
|
|
1c4b51b990 | ||
|
|
2e2c47928b | ||
|
|
80abe0de7d | ||
|
|
a9f7b2d41c | ||
|
|
d14e551a53 | ||
|
|
68567ef2df | ||
|
|
6bc6f2d86d | ||
|
|
1eb2cc961e | ||
|
|
31124749d1 | ||
|
|
9037498c22 | ||
|
|
db32b53e30 | ||
|
|
b529bfd6c5 | ||
|
|
f3df7a7231 | ||
|
|
485bbe1c6f |
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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}"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 | \* | \* |
|
||||
|
||||
@@ -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. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -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 %}
|
||||
@@ -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
|
||||
|
||||
@@ -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() {
|
||||
|
||||
-91
@@ -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 %}
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()];
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
|
||||
@@ -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 "
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
)
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 request’s 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
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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.")
|
||||
@@ -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()
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user