diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index be9f6927b3..94c16ba118 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -210,7 +210,10 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/models/language/pooling commands: - - pytest -v -s models/language/pooling -m 'not core_model' + # Per-test timeout: a hung test fails with a stack dump instead of holding + # the job until Buildkite cancels it, which also blocks the full run's + # coverage collection. A whole shard takes 12-21 minutes. + - pytest -v -s models/language/pooling -m 'not core_model' --timeout=900 --timeout-method=thread mirror: amd: parallelism: 4 diff --git a/cmake/utils.cmake b/cmake/utils.cmake index 6ded936213..83eed340eb 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -425,6 +425,7 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR if("${_base}" IN_LIST _TGT_CUDA_ARCHS) list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}") + list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a" "${_base}f") list(APPEND _CUDA_ARCHS "${_arch}") endif() endif() @@ -464,8 +465,10 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR string(REGEX REPLACE "[af]$" "" _base "${_arch}") if ("${_base}" IN_LIST _SRC_CUDA_ARCHS) list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_arch}") - list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_base}") list(APPEND _CUDA_ARCHS "${_arch}") + if (NOT "${_base}" IN_LIST _TGT_CUDA_ARCHS) + list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_base}") + endif() endif() endif() endforeach() diff --git a/docker/Dockerfile b/docker/Dockerfile index e29142dd96..cc87f9a21f 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -502,7 +502,11 @@ RUN --mount=type=cache,target=/opt/uv/cache \ > /tmp/nccl-override.txt && \ export UV_OVERRIDE=/tmp/nccl-override.txt; \ fi && \ - export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \ + DEEPEP_ARCH_LIST='9.0a 10.0a' && \ + if [ "$(printf '%s\n' 12.9 "$CUDA_VERSION" | sort -V | head -n1)" = 12.9 ]; then \ + DEEPEP_ARCH_LIST="$DEEPEP_ARCH_LIST 10.3a"; \ + fi && \ + export TORCH_CUDA_ARCH_LIST="$DEEPEP_ARCH_LIST" && \ /tmp/install_python_libraries.sh \ --workspace /tmp/ep_kernels_workspace \ --mode wheel \ diff --git a/docs/models/pooling_models/scoring.md b/docs/models/pooling_models/scoring.md index 7ff186b5cd..3bda12a09c 100644 --- a/docs/models/pooling_models/scoring.md +++ b/docs/models/pooling_models/scoring.md @@ -466,6 +466,31 @@ This approach is more robust than index-based access (`messages[0]`, `messages[1 Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../../examples/pooling/score/template/nemotron-rerank.jinja) +#### Templates declared by Sentence Transformers checkpoints + +vLLM does not apply a checkpoint's chat template by default, since many rerankers inherit an unrelated template from their base LLM. The exception is a Sentence Transformers (v5.4+) `CrossEncoder` with a single `sequence-classification` Transformer module whose `sentence_bert_config.json` declares the `message` modality. Sentence Transformers formats the pairs of such checkpoints as `query`/`document` messages, so vLLM does the same with the saved template: + +```json +{ + "transformer_task": "sequence-classification", + "modality_config": { + "text": {"method": "forward", "method_output_name": "logits"}, + "message": {"method": "forward", "method_output_name": "logits", "format": "flat"} + }, + "module_output_name": "scores", + "processing_kwargs": {"chat_template": {"add_generation_prompt": true}} +} +``` + +- The default chat template is used, unless `processing_kwargs.chat_template.chat_template` names another template saved with the tokenizer (e.g. `additional_chat_templates/score.jinja` for `"score"`). +- The other `processing_kwargs.chat_template` settings, such as `add_generation_prompt`, are passed to the template; the request's `chat_template_kwargs` override them. +- Message contents are passed as strings, as with `--chat-template`. +- vLLM does not restore the template suffix after truncation (`restore_suffix`). Use `max_tokens_per_doc` rather than `truncate_prompt_tokens` to truncate long documents. +- `--chat-template` (or the `chat_template` argument of `LLM.score`) takes precedence over the saved template and its settings. +- Models with a built-in score template keep using it. + +The template in use is logged at startup. + ### Enable/disable activation You can enable or disable activation via `use_activation` only works for cross-encoder models. diff --git a/docs/models/pooling_models/specific_models.md b/docs/models/pooling_models/specific_models.md index 0c7b0a1bf3..5b2871ca92 100644 --- a/docs/models/pooling_models/specific_models.md +++ b/docs/models/pooling_models/specific_models.md @@ -313,14 +313,12 @@ Start the server: ```shell vllm serve nvidia/llama-nemotron-rerank-vl-1b-v2 \ --runner pooling \ - --trust-remote-code \ - --chat-template examples/pooling/score/template/nemotron-vl-rerank.jinja + --trust-remote-code ``` !!! note - The chat template bundled with this checkpoint's tokenizer is not suitable - for the Score/Rerank APIs. Use the provided override template when serving: - `examples/pooling/score/template/nemotron-vl-rerank.jinja`. + The checkpoint declares its score template in its Sentence Transformers config, so vLLM applies it automatically (see [Templates declared by Sentence Transformers checkpoints](scoring.md#templates-declared-by-sentence-transformers-checkpoints)). + Revisions older than October 2026 need `--chat-template examples/pooling/score/template/nemotron-vl-rerank.jinja`. Score a text query against an image document: diff --git a/docs/serving/online_serving/token_in_token_out.md b/docs/serving/online_serving/token_in_token_out.md index eade608bc4..6279fbd78b 100644 --- a/docs/serving/online_serving/token_in_token_out.md +++ b/docs/serving/online_serving/token_in_token_out.md @@ -1,8 +1,50 @@ # Tokens In <> Tokens Out API -`/inference/v1/generate` takes prompt token IDs and returns generated token IDs. It is the generation step of the `render → generate → derender` pipeline used in disaggregated serving (see the [Renderer APIs](renderer.md) and [Derenderer APIs](derenderer.md)). +`POST /inference/v1/generate` is the token-in/token-out (TITO) HTTP endpoint: it takes prompt token IDs and returns generated token IDs, logprobs keyed by integer token ID, and `finish_reason`/`stop_reason`, with no detokenization. Use it when the caller does its own tokenization and chat templating, such as an RL rollout worker. It is also the generation step of the `render → generate → derender` pipeline used in disaggregated serving (see the [Renderer APIs](renderer.md) and [Derenderer APIs](derenderer.md)). -The endpoint is registered on `vllm serve` with `--enable-scale-out`, and always with `--tokens-only`. `--tokens-only` also skips loading the tokenizer, so the server can't turn token IDs back into text. +## Enabling the endpoint (`--enable-scale-out`, `--tokens-only`) + +- `--enable-scale-out` registers it, with `/render`, `/derender` and `/inference/v1/abort_requests`. +- `--tokens-only` always registers it, skips loading the tokenizer and forces `detokenize=False`. Stop strings (`sampling_params.stop`) are then rejected with a 400; use `stop_token_ids`. +- With neither flag, it returns 403 `Scale-out endpoints are disabled. Restart vllm serve with --enable-scale-out to enable them.` + +```bash +vllm serve Qwen/Qwen3-0.6B --enable-scale-out + +curl -X POST http://localhost:8000/inference/v1/generate \ + -H "Content-Type: application/json" \ + -d '{"token_ids": [151644, 872, 198], "sampling_params": {"max_tokens": 32, "logprobs": 1}}' +``` + +A full client is in [examples/scale_out/token_generation_client.py](../../../examples/scale_out/token_generation_client.py). + +## Request fields + +| Field | Notes | +| --- | --- | +| `token_ids` | Required `list[int]`: the already-templated prompt. Non-empty, no negative IDs. | +| `sampling_params` | Required. Any [SamplingParams](../../api/README.md#inference-parameters) field, e.g. `max_tokens`, `temperature`, `logprobs`, `stop_token_ids`. Without `max_tokens`, the server applies its own default (as on the OpenAI endpoints), not 16. | +| `output_mode` | `"tokens"` (default) or `"text"`, see [Output modes](#output-modes). | +| `stream`, `stream_options` | SSE chunks ending with `data: [DONE]`. | +| `return_token_ids` | Echo the final prompt IDs as `prompt_token_ids`. | +| `features` / `content_parts` | Multimodal input, mutually exclusive. | +| `request_id`, `model`, `cache_salt`, `priority`, `kv_transfer_params`, `ec_transfer_params` | Optional. | + +## Response fields + +| `choices[i]` field | Notes | +| --- | --- | +| `token_ids` | Generated token IDs. | +| `finish_reason` | `stop`, `length` or `abort`. | +| `stop_reason` | Stop token ID or string; `null` on EOS, length and abort. | +| `logprobs` | `{"content": [{"token_id", "logprob", "rank", "top_logprobs": [...]}]}`, one entry per token, sampled token first. `token_id` is an integer, not the older `"token_id:"` string. Needs `sampling_params.logprobs`. | +| `weight_versions`, `routed_experts`, `sampling_mask` | Optional, see [Weight versions](#weight-versions); the last two need `--enable-return-routed-experts` / `--return-sampling-mask`. | + +The top level adds `request_id`, `usage`, `prompt_token_ids` and `prompt_logprobs`. Stream chunks carry the same per-choice fields for the new tokens. + +## Targeting a data-parallel rank (`X-data-parallel-rank`) + +Send the header `X-data-parallel-rank: ` to pin a request to one data-parallel (DP) engine rank; there is no body field. Without the header, or with a non-integer value, the request is load balanced. A rank outside `[0, data_parallel_size)` is rejected. ## Output modes diff --git a/docs/training/rlhf.md b/docs/training/rlhf.md index 2aae267e1d..ccce0f37dd 100644 --- a/docs/training/rlhf.md +++ b/docs/training/rlhf.md @@ -18,6 +18,8 @@ The following open-source RL libraries use vLLM for fast rollouts (sorted alphab For weight synchronization between training and inference, see the [Weight Transfer](weight_transfer/README.md) documentation, which covers the pluggable backend system with [NCCL](weight_transfer/nccl.md) (multi-GPU) and [IPC](weight_transfer/ipc.md) (same-GPU) engines. To verify that an update landed on every rank, see the [Weight Checker](weight_checker.md). +For MoE policies, see [Routed Experts Return](routed_experts.md) (`--enable-return-routed-experts`) to get the per-token expert routing of each rollout for routing replay (R3) in the trainer. + For pipelining generation and training to improve GPU utilization and throughput, see the [Async Reinforcement Learning](async_rl.md) guide, which covers the pause/resume API for safely updating weights mid-flight. See the following notebooks showing how to use vLLM for GRPO: diff --git a/docs/training/routed_experts.md b/docs/training/routed_experts.md new file mode 100644 index 0000000000..f531a0ebbd --- /dev/null +++ b/docs/training/routed_experts.md @@ -0,0 +1,54 @@ +# Routed Experts Return (MoE Routing Replay / R3) + +`--enable-return-routed-experts` makes vLLM return the routed experts (the expert IDs each MoE layer's router selected) for every token of a request. RL trainers use it for MoE routing replay (R3): the training forward pass reuses the rollout's routing. + +```bash +vllm serve Qwen/Qwen3-30B-A3B --enable-return-routed-experts +# offline: LLM("Qwen/Qwen3-30B-A3B", enable_return_routed_experts=True) +``` + +The flag sets `AuxOutputConfig.enable_return_routed_experts` (full form: `--aux-output-config '{"enable_return_routed_experts": true, "max_bytes": N}'`). There is no per-request on/off switch. + +## Requirements + +Startup fails with a `ValueError` unless: + +- **Model Runner V2** is used (`VLLM_USE_V2_MODEL_RUNNER=1` if vLLM falls back to V1); +- the runner is `generate` and the model is **MoE**; +- **prefix caching is enabled** (the default). `--no-enable-prefix-caching` fails; the old v0.21 advice to disable it is wrong. + +Rejected at startup: adaptive speculative verification, PP > 1, DCP/PCP > 1, and the `NixlConnector`, `NixlPullConnector`, `NixlPushConnector`, `MoRIIOConnector`, `MooncakeConnector` KV connectors. Regular speculative decoding works. + +### How it works: aux output connector + +Routing is stored by the aux output connector (`vllm/distributed/aux_output_connector/`, configured via `--aux-output-config`), keyed by KV block hash and shared with prefix caching, which is why prefix caching is required. `max_bytes` is its host-memory LRU capacity (default: derived from KV cache capacity). If a prefix-cache hit needs evicted rows, the worker raises `BlockObjectStoreError`; increase `max_bytes`. + +## Output format + +| API | Field | +| --- | --- | +| Offline (`LLM`, `AsyncLLM`) | `CompletionOutput.routed_experts`: `np.ndarray` or `None` | +| `/v1/completions`, `/v1/chat/completions` | `choices[i].routed_experts`: base64 `.npy`, **non-streaming only** | +| `/inference/v1/generate` | `routed_experts` per choice (base64 `.npy`); with `stream: true`, on the final chunk | + +- **Shape:** `(num_tokens - 1 - routed_experts_prompt_start, num_moe_layers, num_experts_per_tok)`, `num_tokens` = prompt + output. The last sampled token has no row. +- **Rows:** row `i` is position `routed_experts_prompt_start + i` of `prompt_token_ids + output_token_ids`. +- **dtype:** `uint8` for at most 256 experts, else `uint16`; logical expert IDs. +- Filled only when the request finishes. `None` if the flag is off or the request aborted before any forward pass. + +```python +import base64, io +import numpy as np + +experts = np.load(io.BytesIO(base64.b64decode(choice["routed_experts"]))) +``` + +## `routed_experts_prompt_start` + +Per-request, default `0`: drops the first N prompt-token rows (e.g. routing already received in earlier turns). Must be `0 <= value <= prompt_len`, else HTTP 400. Pass it top-level on `/v1/*completions`, inside `sampling_params` on `/inference/v1/generate`, or as `SamplingParams(routed_experts_prompt_start=...)`. + +## Limitations + +- The Rust HTTP frontend never returns `routed_experts` and always sends `routed_experts_prompt_start = 0`. Use the Python API server. + +Example: [examples/rl/routed_experts_e2e.py](../../examples/rl/routed_experts_e2e.py). diff --git a/docs/training/sampling_mask.md b/docs/training/sampling_mask.md index 7de2366800..25dfb414d0 100644 --- a/docs/training/sampling_mask.md +++ b/docs/training/sampling_mask.md @@ -139,8 +139,9 @@ Each position then carries the sampled token's logprob and the top-128 drop them. The finite entries are the sampling mask when it has at most `k` tokens. - The sampled token is always the first entry and may repeat inside the top-k. -- `/inference/v1/generate` encodes IDs as `token_id:`; the OpenAI - endpoints need `return_tokens_as_token_ids`. +- The `token_id:` strings come from the OpenAI endpoints with + `return_tokens_as_token_ids`; `/inference/v1/generate` returns an integer + `token_id` instead. ## Limitations diff --git a/docs/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md index c8382e2a06..98666e7b28 100644 --- a/docs/training/weight_transfer/README.md +++ b/docs/training/weight_transfer/README.md @@ -145,7 +145,7 @@ engine.send_weights() # called on every rank ## API Endpoints -When running vLLM as an HTTP server, the following endpoints are available for weight transfer. `HTTPVLLMWeightSyncClient` speaks the first four for you. +When running vLLM as an HTTP server, the following endpoints are available for weight transfer. `HTTPVLLMWeightSyncClient` speaks the first five for you. A trainer may also implement the protocol directly without importing vLLM: see the wire contracts in [NCCL](nccl.md#trainer-without-vllm-over-http) and [IPC](ipc.md#trainer-without-vllm-over-http). This keeps the trainer decoupled from vLLM: it needs no vLLM install or version pin, only this documented HTTP and wire contract, so trainer and server can upgrade independently (non-PyTorch trainers such as JAX can use it too). | Endpoint | Method | Description | | -------- | ------ | ----------- | @@ -153,6 +153,7 @@ When running vLLM as an HTTP server, the following endpoints are available for w | `/start_weight_update` | POST | Start a weight update | | `/update_weights` | POST | Transfer a batch of weights with backend-specific metadata | | `/finish_weight_update` | POST | Finish the update and optionally commit its `weight_version` | +| `/shutdown_weight_transfer_engine` | POST | Release the transfer channel without stopping inference | | `/update_weight_version` | POST | Update `weight_version` without changing model weights | | `/weight_info` | GET | Get the latest committed weight version | | `/pause` | POST | Pause generation before weight sync to handle inflight requests | @@ -162,6 +163,12 @@ When running vLLM as an HTTP server, the following endpoints are available for w !!! note The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set. +!!! warning + If `/update_weights` fails, the model is left mid-reload while `/health` + stays 200, and `/finish_weight_update` is rejected. Before generating again, + resend a full `/start_weight_update` → `/update_weights` → + `/finish_weight_update` (see [#57854](https://github.com/vllm-project/vllm/pull/57854)). + The Rust frontend's optional gRPC `Control` service exposes the same pause, sleep, weight-transfer, and weight-version lifecycle for trusted sidecars. The `ServerInfo.rl_capabilities` response reports whether weight transfer and sleep mode were configured. Backend-specific `init_info` and `update_info` remain JSON metadata; model tensors continue to move over the configured NCCL, IPC, sparse-NCCL, or sharded-RDT transport. ## Extending the System diff --git a/docs/training/weight_transfer/ipc.md b/docs/training/weight_transfer/ipc.md index b71568c3df..7dab95d66d 100644 --- a/docs/training/weight_transfer/ipc.md +++ b/docs/training/weight_transfer/ipc.md @@ -110,6 +110,43 @@ mode carries a per-chunk barrier across ranks: without it, a rank could overwrit its buffer while its colocated worker is still reading the current chunk. This is handled inside `send_weights()`. +## Trainer Without vLLM over HTTP + +A colocated trainer can speak the [HTTP endpoints](README.md#api-endpoints) +directly, with only PyTorch and an HTTP client. The server needs +`VLLM_ALLOW_INSECURE_SERIALIZATION=1`; without it, init and start return 200 and +the first `/update_weights` returns 500. Reference: +[`rlhf_http_ipc_no_vllm.py`](../../../examples/rl/rlhf_http_ipc_no_vllm.py). + +`POST /init_weight_transfer_engine` body: `{"init_info": {"packed": false}}`. + +`POST /update_weights` body `{"update_info": {...}}`: + +| Field | Meaning | +| ----- | ------- | +| `names`, `dtype_names`, `shapes` | Checkpoint names, `torch` dtype attribute names (`"bfloat16"`), full shapes | +| `ipc_handles_pickled` | `base64(pickle.dumps(handles))`. Unpacked: a list with one dict per name. Packed: one dict for the shared buffer | +| `tensor_sizes` | Packed only: byte size of each tensor in the chunk, in `names` order | + +- A handle dict is `{gpu_uuid: reduce_tensor(t)[1]}` with one entry per GPU the + server uses (`torch.multiprocessing.reductions.reduce_tensor`). + `gpu_uuid = str(torch.cuda.get_device_properties(i).uuid)`, **without** the + `GPU-` prefix that `nvidia-smi` prints. The worker overwrites `args[6]` + (device index) with its own. +- `ipc_handles` (unpickled) is not JSON-serializable; over HTTP use + `ipc_handles_pickled`. +- Packed: export one `uint8` buffer once and reuse it for every chunk. Copy + `t.contiguous().view(-1).view(torch.uint8)` in `names` order and close a chunk + **before** the next tensor would overflow it (unlike NCCL, which closes after). + A tensor larger than the buffer is an error. Synchronize before each + `/update_weights`; do not overwrite the buffer until the call returns. +- Keep every exported tensor alive until `/finish_weight_update` returns, then + drop the references and call `torch.cuda.ipc_collect()`. +- Multi-rank trainers: each rank exports handles for its own GPU, the ranks + gather the `{uuid: args}` dicts to rank 0 over their own `torch.distributed` + group and merge them per index, and rank 0 sends all HTTP. In packed mode, + barrier before and after each `/update_weights`. + ## Rank-Local Updates For different parameter subsets per worker, `update_info` may be a list indexed @@ -119,4 +156,5 @@ worker to participate. ## Examples - [RLHF with IPC weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_ipc.py) - **Start here.** Server and training model share a single GPU; HTTP control plane, CUDA IPC data plane. Launches and tears down its own server +- [IPC weight syncing from a trainer that does not import vLLM](../../../examples/rl/rlhf_http_ipc_no_vllm.py) - `torch` + `requests` only, `ipc_handles_pickled` over HTTP; see [Trainer Without vLLM over HTTP](#trainer-without-vllm-over-http) - [RLHF with IPC + FSDP2 and expert parallelism](../../../examples/rl/rlhf_ipc_fsdp_ep.py) - Multi-rank trainer colocated with a `--data-parallel-size 4` server on the same 4 GPUs: every FSDP rank builds an engine and joins the handle all-gather, with packed chunking and sleep/wake around the transfer diff --git a/docs/training/weight_transfer/nccl.md b/docs/training/weight_transfer/nccl.md index 3a3e207fdc..f7c73c1ea4 100644 --- a/docs/training/weight_transfer/nccl.md +++ b/docs/training/weight_transfer/nccl.md @@ -130,6 +130,51 @@ pass — this is the invariant to test first. IPC does not read `metadata()` at all: it derives the update info from iteration as it goes, so it cannot observe a divergence. +## Trainer Without vLLM over HTTP + +A trainer can speak the [HTTP endpoints](README.md#api-endpoints) directly, with +only PyTorch and an HTTP client. It mints the `ncclUniqueId` itself and is NCCL +rank 0. Reference: [`rlhf_http_nccl_no_vllm.py`](../../../examples/rl/rlhf_http_nccl_no_vllm.py). + +`POST /init_weight_transfer_engine` body `{"init_info": {...}}`: + +| Field | Default | Meaning | +| ----- | ------- | ------- | +| `nccl_unique_id_b64` | — | Standard (not URL-safe) base64 of the 128-byte `ncclUniqueId`. Replaces `master_address` / `master_port` | +| `rank_offset` | — | First worker rank; `1` for one trainer rank | +| `world_size` | — | Trainer ranks + `GET /get_world_size` | +| `packed` | `false` | Packed broadcasting | +| `packed_buffer_size_bytes` | 1 GiB | Must match the trainer's chunking | +| `packed_num_buffers` | 2 | Worker-side receive buffers | + +- Worker rank = `rank_offset + dp_index * (TP * PP) + rank`. +- `/init_weight_transfer_engine` returns only after all ranks join: send it in + the background while the trainer calls `ncclCommInitRank(world_size, uid, 0)`. + There is no timeout; if the trainer never joins, the workers wait forever. +- Right after init every worker issues a warm-up `all_reduce` (sum) of + `torch.zeros(1)` (one float32 element). The trainer must issue the same one + before any other collective. + +Each sync: `POST /start_weight_update`, then `POST /update_weights` with +`{"update_info": {"names", "dtype_names", "shapes"}}` **in the background** +while rank 0 broadcasts in `names` order, then `POST /finish_weight_update`. +`dtype_names` are `torch` attribute names (`"bfloat16"`); `shapes` are full +checkpoint shapes. + +- Unpacked: one broadcast per tensor. +- Packed: concatenate `t.contiguous().view(-1).view(torch.uint8)` in `names` + order; broadcast a chunk as soon as its size **exceeds** + `packed_buffer_size_bytes` (the crossing tensor stays in the chunk), then the + remainder. + +Trainer-side communicators without vLLM: + +| Communicator | Status | +| ------------ | ------ | +| `torch.distributed.ProcessGroupNCCL` on a private `TCPStore` + `eager_connect_single_device` | Works; the unique id is read from store key `"0"`, a torch implementation detail (used by the example) | +| `ctypes` against torch's bundled `libnccl.so.2` | Works | +| `torch.cuda.nccl.init_rank` | Fails on Python 3.10+ (`PY_SSIZE_T_CLEAN` `SystemError`) | + ## Sparse NCCL Sparse, flat-index weight patches use @@ -196,6 +241,7 @@ Sparse NCCL sends only `O(nnz)` indices and values. Checkpoint application still ## Examples - [RLHF with NCCL weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - **Start here.** Trainer on one GPU, 2x tensor-parallel fp8 server on two others; HTTP control plane, NCCL data plane. Launches and tears down its own server +- [NCCL weight syncing from a trainer that does not import vLLM](../../../examples/rl/rlhf_http_nccl_no_vllm.py) - `torch` + `requests` only, `nccl_unique_id_b64` rendezvous; see [Trainer Without vLLM over HTTP](#trainer-without-vllm-over-http) - [RLHF with NCCL + FSDP2 and expert parallelism](../../../examples/rl/rlhf_nccl_fsdp_ep.py) - Multi-rank trainer: every FSDP rank builds an engine and joins the `full_tensor()` gather while only rank 0 touches the wire - [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Qwen3 MoE per-expert checkpoint updates with one trainer GPU and two TP2/EP2 inference GPUs - [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model; uses an in-process `AsyncLLMEngine` with `RayVLLMWeightSyncClient` diff --git a/examples/rl/rlhf_http_ipc_no_vllm.py b/examples/rl/rlhf_http_ipc_no_vllm.py new file mode 100644 index 0000000000..20a92a149c --- /dev/null +++ b/examples/rl/rlhf_http_ipc_no_vllm.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CUDA IPC weight sync over HTTP from a trainer that never imports vLLM. + +The trainer uses only `torch`, `requests`, `safetensors`, `pybase64` and the +stdlib, and speaks the wire protocol in +`docs/training/weight_transfer/ipc.md` ("Trainer Without vLLM over HTTP"): +it exports its tensors with `torch.multiprocessing.reductions.reduce_tensor` +and sends the handles, pickled and base64-encoded, in `ipc_handles_pickled`. + +IPC requires colocation: this process holds a full copy of the weights on every +GPU the server uses, and keys each handle by that GPU's UUID. A multi-process +trainer would instead gather the per-rank `{uuid: args}` dicts to one rank. + +Start a server; IPC handles are pickled, so it needs +`VLLM_ALLOW_INSECURE_SERIALIZATION=1` (the weights it starts with are discarded): + + VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_INSECURE_SERIALIZATION=1 \ + CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-0.6B \ + --load-format dummy --enforce-eager --gpu-memory-utilization 0.4 \ + --weight-transfer-config '{"backend": "ipc"}' + +Then run the trainer on the same GPU(s) (`--packed` for chunked transfer): + + CUDA_VISIBLE_DEVICES=0 python examples/rl/rlhf_http_ipc_no_vllm.py \ + --checkpoint /path/to/Qwen3-0.6B + +For a TP=2 server on GPUs 0-1, use `CUDA_VISIBLE_DEVICES=0,1` on both sides and +pass `--num-gpus 2` to the trainer. +""" + +import sys + + +class _BlockVLLM: + def find_spec(self, name, path=None, target=None): + if name.split(".")[0] == "vllm": + raise ImportError(f"this trainer must not import {name!r}") + return None + + +sys.meta_path.insert(0, _BlockVLLM()) + +import argparse # noqa: E402 +import glob # noqa: E402 +import pickle # noqa: E402 + +import pybase64 as base64 # noqa: E402 +import requests # noqa: E402 +import torch # noqa: E402 +from safetensors.torch import load_file # noqa: E402 +from torch.multiprocessing.reductions import reduce_tensor # noqa: E402 + +parser = argparse.ArgumentParser() +parser.add_argument("--url", default="http://localhost:8000") +parser.add_argument("--checkpoint", required=True, help="dir with *.safetensors") +parser.add_argument("--num-gpus", type=int, default=1) +parser.add_argument("--packed", action="store_true") +parser.add_argument("--packed-buffer-size-bytes", type=int, default=512 << 20) +args = parser.parse_args() + + +def post(path: str, body: dict | None = None) -> dict: + r = requests.post(args.url + path, json=body or {}, timeout=600) + r.raise_for_status() + return r.json() + + +def generate(prompts: list[str]) -> list[str]: + model = requests.get(args.url + "/v1/models", timeout=60).json()["data"][0]["id"] + body = {"model": model, "prompt": prompts, "max_tokens": 16, "temperature": 0} + choices = post("/v1/completions", body)["choices"] + return [c["text"] for c in sorted(choices, key=lambda c: c["index"])] + + +# Weights are {gpu_uuid: {name: tensor}}: one full copy per colocated GPU. The +# UUID is str(torch.cuda.get_device_properties(i).uuid), with no "GPU-" prefix. +Weights = dict[str, dict[str, torch.Tensor]] + + +def update_weights( + weights: Weights, + names: list[str], + ipc_handles: list[dict[str, tuple]] | dict[str, tuple], + tensor_sizes: list[int] | None = None, +) -> None: + w = next(iter(weights.values())) + update_info = { + "names": names, + "dtype_names": [str(w[n].dtype).removeprefix("torch.") for n in names], + "shapes": [list(w[n].shape) for n in names], + "ipc_handles_pickled": base64.b64encode(pickle.dumps(ipc_handles)).decode(), + } + if tensor_sizes is not None: + update_info["tensor_sizes"] = tensor_sizes + post("/update_weights", {"update_info": update_info}) + + +def send_unpacked(weights: Weights, names: list[str]) -> list[torch.Tensor]: + """One /update_weights call; one {uuid: args} dict per parameter.""" + copies = {u: {n: w[n].contiguous() for n in names} for u, w in weights.items()} + handles = [{u: reduce_tensor(copies[u][n])[1] for u in copies} for n in names] + update_weights(weights, names, handles) + return [t for c in copies.values() for t in c.values()] + + +def send_packed(weights: Weights, names: list[str], cap: int) -> list[torch.Tensor]: + """One /update_weights call per chunk; one {uuid: args} dict per chunk.""" + buffers = { + u: torch.empty(cap, dtype=torch.uint8, device=w[names[0]].device) + for u, w in weights.items() + } + handle = {u: reduce_tensor(b)[1] for u, b in buffers.items()} # export once + chunk: list[str] = [] + sizes: list[int] = [] + + def flush() -> None: + for b in buffers.values(): + torch.accelerator.synchronize(b.device) # copies must land before the read + update_weights(weights, chunk, handle, sizes) + chunk.clear() + sizes.clear() + + for n in names: + nbytes = next(iter(weights.values()))[n].nbytes + assert nbytes <= cap, f"{n} ({nbytes} B) exceeds the buffer" + if chunk and sum(sizes) + nbytes > cap: # close *before* overflowing + flush() + offset = sum(sizes) + for u, w in weights.items(): + flat = w[n].contiguous().view(-1).view(torch.uint8) + buffers[u][offset : offset + nbytes].copy_(flat) + chunk.append(n) + sizes.append(nbytes) + if chunk: + flush() + return list(buffers.values()) + + +def main() -> None: + weights: Weights = {} + for i in range(args.num_gpus): + uuid = str(torch.cuda.get_device_properties(i).uuid) + weights[uuid] = {} + for path in sorted(glob.glob(f"{args.checkpoint}/*.safetensors")): + weights[uuid].update(load_file(path, device=f"cuda:{i}")) + names = sorted(next(iter(weights.values()))) + + prompts = ["The capital of France is", "def fibonacci(n):"] + print("before:", generate(prompts)) + + post("/init_weight_transfer_engine", {"init_info": {"packed": args.packed}}) + post("/start_weight_update") + if args.packed: + exported = send_packed(weights, names, args.packed_buffer_size_bytes) + else: + exported = send_unpacked(weights, names) + post("/finish_weight_update", {"weight_version": None}) + # The worker reads the exported memory until finish; only then release it. + del exported + torch.cuda.ipc_collect() + + print("after: ", generate(prompts)) + assert not [m for m in sys.modules if m.split(".")[0] == "vllm"] + + +if __name__ == "__main__": + main() diff --git a/examples/rl/rlhf_http_nccl_no_vllm.py b/examples/rl/rlhf_http_nccl_no_vllm.py new file mode 100644 index 0000000000..c9d8abad36 --- /dev/null +++ b/examples/rl/rlhf_http_nccl_no_vllm.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NCCL weight sync over HTTP from a trainer that never imports vLLM. + +The trainer uses only `torch`, `requests`, `safetensors`, `pybase64` and the +stdlib, and speaks the wire protocol in +`docs/training/weight_transfer/nccl.md` ("Trainer Without vLLM over HTTP"): +`nccl_unique_id_b64` rendezvous, a one-element warm-up `all_reduce`, then +`/start_weight_update` -> `/update_weights` (concurrent with the rank-0 +broadcast) -> `/finish_weight_update`. + +The NCCL communicator is a plain `torch.distributed.ProcessGroupNCCL` over a +private `TCPStore` that only the trainer joins. Its `ncclUniqueId` is read back +from store key "0", which is a torch implementation detail. + +Start a server on GPUs 0-1 (the weights it starts with are discarded): + + VLLM_SERVER_DEV_MODE=1 CUDA_VISIBLE_DEVICES=0,1 vllm serve Qwen/Qwen3-0.6B \ + --tensor-parallel-size 2 --load-format dummy --enforce-eager \ + --weight-transfer-config '{"backend": "nccl"}' + +Then run the trainer on GPU 2 (`--packed` for packed broadcasting): + + CUDA_VISIBLE_DEVICES=2 python examples/rl/rlhf_http_nccl_no_vllm.py \ + --checkpoint /path/to/Qwen3-0.6B +""" + +import sys + + +class _BlockVLLM: + def find_spec(self, name, path=None, target=None): + if name.split(".")[0] == "vllm": + raise ImportError(f"this trainer must not import {name!r}") + return None + + +sys.meta_path.insert(0, _BlockVLLM()) + +import argparse # noqa: E402 +import datetime # noqa: E402 +import glob # noqa: E402 +import threading # noqa: E402 + +import pybase64 as base64 # noqa: E402 +import requests # noqa: E402 +import torch # noqa: E402 +import torch.distributed as dist # noqa: E402 +from safetensors.torch import load_file # noqa: E402 + +parser = argparse.ArgumentParser() +parser.add_argument("--url", default="http://localhost:8000") +parser.add_argument("--checkpoint", required=True, help="dir with *.safetensors") +parser.add_argument("--packed", action="store_true") +parser.add_argument("--packed-buffer-size-bytes", type=int, default=64 << 20) +parser.add_argument("--store-port", type=int, default=29612) +args = parser.parse_args() + + +def post(path: str, body: dict | None = None) -> dict: + r = requests.post(args.url + path, json=body or {}, timeout=600) + r.raise_for_status() + return r.json() + + +class PostInBackground(threading.Thread): + """An HTTP call that only returns once this process joins the collective.""" + + def __init__(self, path: str, body: dict) -> None: + super().__init__() + self.path, self.body = path, body + self.error: BaseException | None = None + self.start() + + def run(self) -> None: + try: + post(self.path, self.body) + except BaseException as e: + self.error = e + + def result(self) -> None: + self.join() + if self.error is not None: + raise self.error + + +def generate(prompts: list[str]) -> list[str]: + model = requests.get(args.url + "/v1/models", timeout=60).json()["data"][0]["id"] + body = {"model": model, "prompt": prompts, "max_tokens": 16, "temperature": 0} + choices = post("/v1/completions", body)["choices"] + return [c["text"] for c in sorted(choices, key=lambda c: c["index"])] + + +def main() -> None: + device = torch.device("cuda:0") + torch.accelerator.set_device_index(device) + weights: dict[str, torch.Tensor] = {} + for path in sorted(glob.glob(f"{args.checkpoint}/*.safetensors")): + weights.update(load_file(path, device="cuda:0")) + names = sorted(weights) + + prompts = ["The capital of France is", "def fibonacci(n):"] + print("before:", generate(prompts)) + + # Trainer is NCCL rank 0; vLLM worker i is rank 1 + i. + world_size = 1 + requests.get(args.url + "/get_world_size").json()["world_size"] + store = dist.TCPStore( + "127.0.0.1", args.store_port, is_master=True, wait_for_workers=False + ) + pg = dist.ProcessGroupNCCL(store, 0, world_size, dist.ProcessGroupNCCL.Options()) + connect = threading.Thread(target=pg.eager_connect_single_device, args=(device,)) + connect.start() # publishes the unique id, then waits in ncclCommInitRank + reader = dist.TCPStore("127.0.0.1", args.store_port, is_master=False) + reader.wait(["0"], datetime.timedelta(seconds=60)) + unique_id = reader.get("0") + assert len(unique_id) == 128 + + init = PostInBackground( + "/init_weight_transfer_engine", + { + "init_info": { + "nccl_unique_id_b64": base64.b64encode(unique_id).decode(), + "rank_offset": 1, + "world_size": world_size, + "packed": args.packed, + "packed_buffer_size_bytes": args.packed_buffer_size_bytes, + } + }, + ) + connect.join() + pg.allreduce([torch.zeros(1, device=device)]).wait() # vLLM's warm-up + init.result() + + def broadcast(t: torch.Tensor) -> None: + opts = dist.BroadcastOptions() + opts.rootRank = 0 + pg.broadcast([t], opts).wait() + + post("/start_weight_update") + update_info = { + "names": names, + "dtype_names": [str(weights[n].dtype).removeprefix("torch.") for n in names], + "shapes": [list(weights[n].shape) for n in names], + } + update = PostInBackground("/update_weights", {"update_info": update_info}) + if not args.packed: + for n in names: + broadcast(weights[n].contiguous()) + else: + # Same chunking as the worker: a chunk closes right after the tensor + # that pushes it past the buffer size. + chunk: list[torch.Tensor] = [] + size = 0 + for n in names: + chunk.append(weights[n].contiguous().view(-1).view(torch.uint8)) + size += chunk[-1].numel() + if size > args.packed_buffer_size_bytes: + broadcast(torch.cat(chunk)) + chunk, size = [], 0 + if chunk: + broadcast(torch.cat(chunk)) + torch.accelerator.synchronize() + update.result() + post("/finish_weight_update", {"weight_version": None}) + + print("after: ", generate(prompts)) + post("/shutdown_weight_transfer_engine") + pg.shutdown() + + assert not [m for m in sys.modules if m.split(".")[0] == "vllm"] + + +if __name__ == "__main__": + main() diff --git a/examples/scale_out/token_generation_client.py b/examples/scale_out/token_generation_client.py index a4792a68f9..063e8ae39b 100644 --- a/examples/scale_out/token_generation_client.py +++ b/examples/scale_out/token_generation_client.py @@ -3,6 +3,8 @@ """Call a token-in/token-out server started with: vllm serve Qwen/Qwen3-0.6B --enable-scale-out + +Fields and DP rank routing: docs/serving/online_serving/token_in_token_out.md """ import httpx diff --git a/mkdocs.yaml b/mkdocs.yaml index bed8318998..d591df7b14 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -106,6 +106,7 @@ plugins: examples/features/pause_resume.md: examples/rl/pause_resume.md serving/online_serving/decisions.md: models/structured_decisions_models.md serving/online_serving/structured_decisions.md: models/structured_decisions_models.md + training/routed_experts_replay.md: training/routed_experts.md markdown_extensions: - attr_list diff --git a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py index 7f2c157f25..d5a585f579 100644 --- a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py +++ b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py @@ -65,6 +65,8 @@ def llm_reranker_processor() -> CrossEncoderIOProcessor: processor.supports_score_template = False processor.use_sep_token = False processor.model = None + processor.saved_chat_template = None + processor.saved_chat_template_kwargs = {} return processor diff --git a/tests/models/language/pooling_mteb_test/test_nemotron.py b/tests/models/language/pooling_mteb_test/test_nemotron.py index 711224a1a2..3c3618de65 100644 --- a/tests/models/language/pooling_mteb_test/test_nemotron.py +++ b/tests/models/language/pooling_mteb_test/test_nemotron.py @@ -39,6 +39,14 @@ RERANK_MODELS = [ is_prefix_caching_supported=False, is_chunked_prefill_supported=False, ), + # Declares its score template in its Sentence Transformers config, so no + # chat template is passed. + RerankModelInfo( + "nvidia/llama-nemotron-rerank-vl-1b-v2", + architecture="LlamaNemotronVLForSequenceClassification", + mteb_score=0.33919, + mteb_tol=1e-2, + ), ] diff --git a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py index 3cbeda2e49..227542023a 100644 --- a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py +++ b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py @@ -10,7 +10,6 @@ Both variants share a SigLIP vision encoder with a bidirectional LLaMA backbone. from collections.abc import Sequence from io import BytesIO -from pathlib import Path import pybase64 as base64 import pytest @@ -202,13 +201,6 @@ def test_models( RERANKER_MODELS = ["nvidia/llama-nemotron-rerank-vl-1b-v2"] -# The tokenizer's built-in chat template is not suitable for the Score/Rerank -# APIs (it's inherited from the base LLM). We must use the provided override. -_RERANKER_SCORE_TEMPLATE = ( - Path(__file__).parents[4] - / "examples/pooling/score/template/nemotron-vl-rerank.jinja" -).read_text() - RERANKER_TEXT_QUERY = "How is AI improving the intelligence and capabilities of robots?" RERANKER_TEXT_DOCS = [ "AI enables robots to perceive, plan, and act autonomously.", @@ -281,7 +273,11 @@ def _run_vllm_reranker( dtype: str, input_cases: Sequence[RerankerCase], ) -> list[list[float]]: - """Run all vLLM reranker cases in one model lifecycle.""" + """Run all vLLM reranker cases in one model lifecycle. + + The checkpoint declares its score template in its Sentence Transformers + config, so no chat template is passed. + """ with vllm_runner( model, runner="pooling", @@ -298,11 +294,7 @@ def _run_vllm_reranker( if not has_images: queries = [query] * len(docs) doc_texts = [doc_text for doc_text, _ in docs] - outputs = vllm_model.score( - queries, - doc_texts, - chat_template=_RERANKER_SCORE_TEMPLATE, - ) + outputs = vllm_model.score(queries, doc_texts) else: query_params = [ ScoreMultiModalParam( @@ -334,11 +326,7 @@ def _run_vllm_reranker( ) doc_params.append(ScoreMultiModalParam(content=content)) - raw_outputs = vllm_model.llm.score( - query_params, - doc_params, - chat_template=_RERANKER_SCORE_TEMPLATE, - ) + raw_outputs = vllm_model.llm.score(query_params, doc_params) outputs = [output.outputs.score for output in raw_outputs] scores_per_case.append(outputs) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 4380f934d2..16ed5c95cf 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -1,16 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging import mmap import os from dataclasses import dataclass -from functools import partial from itertools import accumulate from pathlib import Path from types import SimpleNamespace import pytest import torch +from compressed_tensors.quantization import QuantizationArgs from safetensors import safe_open from safetensors.torch import load, load_file, save_file from torch import nn @@ -25,6 +26,10 @@ from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata from vllm.config import SpeculativeConfig, VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.config.quantization import QuantizationConfigArgs +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 + CompressedTensorsConfig, +) from vllm.model_executor.layers.quantization.fp8 import Fp8Config from vllm.model_executor.layers.quantization.inc import INCConfig from vllm.model_executor.layers.quantization.modelopt import ( @@ -39,8 +44,7 @@ from vllm.models.qwen4_exp.amd.ple_layer import ( ) from vllm.models.qwen4_exp.common.ple import ( PLEShardOverlap, - compute_ple_shard_overlap, - copy_ple_embedding_shard_, + PLEVocabParallelEmbedding, ) from vllm.models.qwen4_exp.nvidia.ngram_embedding import ( Qwen4ExpNGramEmbedding, @@ -217,12 +221,23 @@ def _make_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding: return module +def _make_test_ple_embedding(tp_start: int, tp_end: int) -> PLEVocabParallelEmbedding: + """Create an embedding with only the state needed for shard loading.""" + embedding = PLEVocabParallelEmbedding.__new__(PLEVocabParallelEmbedding) + nn.Module.__init__(embedding) + embedding.shard_indices = SimpleNamespace( + org_vocab_start_index=tp_start, + org_vocab_end_index=tp_end, + ) + return embedding + + def _set_test_embedding_weight_loader(embedding) -> None: - embedding.weight.weight_loader = partial( - copy_ple_embedding_shard_, + loader = _make_test_ple_embedding( tp_start=embedding.shard_indices.org_vocab_start_index, tp_end=embedding.shard_indices.org_vocab_end_index, ) + embedding.weight.weight_loader = loader.weight_loader def _make_fp8_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding: @@ -251,19 +266,16 @@ def _make_fp8_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding: def test_ple_shard_overlap_and_copy() -> None: - overlap = compute_ple_shard_overlap( - checkpoint_start=2, checkpoint_rows=5, tp_start=4, tp_end=8 - ) + embedding = _make_test_ple_embedding(tp_start=4, tp_end=8) + overlap = embedding._compute_shard_overlap(checkpoint_start=2, checkpoint_rows=5) assert overlap == PLEShardOverlap(source_start=2, destination_start=0, row_count=3) destination = torch.full((4, 2), -1.0) loaded = torch.arange(10, dtype=torch.float64).reshape(5, 2) - copied = copy_ple_embedding_shard_( + copied = embedding._copy_embedding_shard_( destination, loaded, checkpoint_start=2, - tp_start=4, - tp_end=8, ) assert copied == 3 @@ -272,13 +284,12 @@ def test_ple_shard_overlap_and_copy() -> None: def test_ple_shard_copy_is_a_noop_without_overlap() -> None: + embedding = _make_test_ple_embedding(tp_start=4, tp_end=8) destination = torch.ones(4, 2) - copied = copy_ple_embedding_shard_( + copied = embedding._copy_embedding_shard_( destination, torch.zeros(2, 2), checkpoint_start=10, - tp_start=4, - tp_end=8, ) assert copied == 0 @@ -711,17 +722,44 @@ def test_ple_fp8_embedding_respects_checkpoint_shard_exclusions() -> None: ) -def test_ple_embedding_rejects_unsupported_quantization_configs() -> None: +@pytest.mark.parametrize( + "quant_config", + [ + pytest.param(ModelOptNvFp4Config(exclude_modules=[]), id="modelopt-nvfp4"), + pytest.param( + OnlineQuantizationConfig(QuantizationConfigArgs(linear="fp8_per_tensor")), + id="online-fp8", + ), + pytest.param( + CompressedTensorsConfig( + target_scheme_map={ + "re:.*ngram_embedding": { + "weights": QuantizationArgs(num_bits=8, type="int"), + }, + }, + ignore=[], + quant_format="int-quantized", + ), + id="compressed-tensors", + ), + ], +) +def test_ple_embedding_falls_back_for_unsupported_quantization_configs( + quant_config: QuantizationConfig, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Warn when global quantization falls back to an unquantized PLE table.""" prefix = "model.layers.1.ple.ple_embedding.ngram_embedding" - nvfp4_config = ModelOptNvFp4Config(exclude_modules=[]) - with pytest.raises(NotImplementedError, match="ModelOptNvFp4Config"): - Qwen4ExpPLEEmbeddingMethod.from_quant_config(nvfp4_config, prefix) - - online_fp8_config = OnlineQuantizationConfig( - QuantizationConfigArgs(linear="fp8_per_tensor") + test_logger = logging.getLogger("test.ple_fallback") + monkeypatch.setattr(ngram_embedding_module, "logger", test_logger) + caplog.set_level(logging.WARNING, logger=test_logger.name) + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix), + Qwen4ExpPLEUnquantizedEmbeddingMethod, ) - with pytest.raises(NotImplementedError, match="OnlineQuantizationConfig"): - Qwen4ExpPLEEmbeddingMethod.from_quant_config(online_fp8_config, prefix) + assert caplog.records[0].levelno == logging.WARNING + assert "Falling back to unquantized PLE" in caplog.text def test_ple_embedding_respects_modelopt_exclusion() -> None: @@ -748,8 +786,10 @@ def test_ple_embedding_respects_inc_layer_config() -> None: ) quant_config.extra_config = None - with pytest.raises(NotImplementedError, match="INCConfig"): - Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix) + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix), + Qwen4ExpPLEUnquantizedEmbeddingMethod, + ) def test_ple_embedding_is_unquantized_under_quark() -> None: @@ -933,12 +973,10 @@ def test_ple_pinned_embedding_loads_on_cpu_and_looks_up_through_uva( loaded_weight = ( torch.arange(12, dtype=torch.float32).reshape(4, 3).to(storage_dtype) ) - copied = copy_ple_embedding_shard_( + copied = embedding._copy_embedding_shard_( embedding.weight, loaded_weight, checkpoint_start=0, - tp_start=0, - tp_end=4, ) input_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0") output = embedding._lookup(input_ids) @@ -995,12 +1033,10 @@ def test_pinned_embedding_start_prefetch_allocates_lazily( loaded_weight = torch.arange(12, dtype=torch.float32).reshape(4, 3) loaded_weight = loaded_weight.to(torch.bfloat16) - copy_ple_embedding_shard_( + embedding._copy_embedding_shard_( embedding.weight, loaded_weight, checkpoint_start=0, - tp_start=0, - tp_end=4, ) ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0") @@ -1019,8 +1055,18 @@ def test_pinned_embedding_start_prefetch_allocates_lazily( torch.testing.assert_close(output, expected, rtol=0, atol=0) -def test_ple_fp8_embedding_supports_mixed_precision_config() -> None: +@pytest.mark.parametrize( + "exclude_pattern", + [ + "model.language_model.layers.1.ple.ple_embedding.ngram_embedding", + "*.ngram_embedding", + ], +) +def test_ple_fp8_embedding_supports_mixed_precision_config( + exclude_pattern: str, +) -> None: prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding" + nvfp4_prefix = "model.language_model.layers.2.ple.ple_embedding.ngram_embedding" quant_config = ModelOptMixedPrecisionConfig.from_config( { "quantization": { @@ -1029,6 +1075,7 @@ def test_ple_fp8_embedding_supports_mixed_precision_config() -> None: "group_size": 16, "quantized_layers": { prefix: {"quant_algo": "FP8"}, + nvfp4_prefix: {"quant_algo": "NVFP4"}, "model.language_model.layers.2.moe.gate_proj": { "quant_algo": "MXFP8" }, @@ -1041,13 +1088,30 @@ def test_ple_fp8_embedding_supports_mixed_precision_config() -> None: Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix), Qwen4ExpPLEFp8EmbeddingMethod, ) - # A quant_algo the PLE table cannot load must not fall back to unquantized. + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config( + quant_config, + "model.language_model.layers.3.ple.ple_embedding.ngram_embedding", + ), + Qwen4ExpPLEUnquantizedEmbeddingMethod, + ) + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, nvfp4_prefix), + Qwen4ExpPLENvFp4EmbeddingMethod, + ) + # An explicitly unsupported PLE format must not fall back to unquantized. with pytest.raises(ValueError, match="Unsupported quant_algo 'MXFP8'"): Qwen4ExpPLEEmbeddingMethod.from_quant_config( quant_config, "model.language_model.layers.2.moe.gate_proj", ) + quant_config.exclude_modules = [exclude_pattern] + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix), + Qwen4ExpPLEUnquantizedEmbeddingMethod, + ) + def _patch_ple_construction( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/models/test_deepseek_v4_aiter_mega_moe.py b/tests/models/test_deepseek_v4_aiter_mega_moe.py index 1a53476694..c3ffded512 100644 --- a/tests/models/test_deepseek_v4_aiter_mega_moe.py +++ b/tests/models/test_deepseek_v4_aiter_mega_moe.py @@ -15,6 +15,7 @@ from vllm.models.deepseek_v4.amd.mega_moe import ( ) from vllm.models.deepseek_v4.amd.model import DeepseekV4MoE from vllm.platforms import current_platform +from vllm.utils.torch_utils import aux_stream HIDDEN = 64 INTER = 64 @@ -156,7 +157,10 @@ def test_shared_expert_route_appended(monkeypatch): @pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm") @pytest.mark.parametrize("hash_routing", [False, True]) @pytest.mark.parametrize("num_tokens", [1, 37]) -def test_mega_moe_routing_matches_reference(monkeypatch, hash_routing, num_tokens): +@pytest.mark.parametrize("shared_stream", [False, True]) +def test_mega_moe_routing_matches_reference( + monkeypatch, hash_routing, num_tokens, shared_stream +): num_experts, top_k, vocab = 16, 4, 64 scaling = 2.5 device = "cuda" @@ -195,6 +199,10 @@ def test_mega_moe_routing_matches_reference(monkeypatch, hash_routing, num_token moe.gate = Gate() moe.experts = experts moe.shared_experts = lambda x: 3 * x + moe.shared_experts_stream = aux_stream() if shared_stream else None + moe.shared_experts_events = ( + (torch.cuda.Event(), torch.cuda.Event()) if shared_stream else None + ) moe.scoring_func = "sqrtsoftplus" moe.n_activated_experts = top_k moe.renormalize = True diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 1c7730a37e..16334fdb0d 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -49,7 +49,9 @@ from vllm.transformers_utils.processor import get_video_processor_cls_name_from_ from .utils import ( create_edit_list_trimmed_video, + create_live_webm_video, create_long_gop_video, + create_raw_h264_video, create_video_from_image, ) @@ -1261,8 +1263,8 @@ def test_qwen3_vl_unknown_source_fps_matches_hf( the metadata carries the fps HF assumes, which the frame timestamps use.""" get_video_metadata = OpenCVVideoBackendMixin.get_video_metadata - def unknown_fps(cap) -> VideoSourceMetadata: - return get_video_metadata(cap)._replace(original_fps=0.0, duration=0.0) + def unknown_fps(data) -> VideoSourceMetadata: + return get_video_metadata(data)._replace(original_fps=0.0, duration=0.0) monkeypatch.setattr( OpenCVVideoBackendMixin, "get_video_metadata", staticmethod(unknown_fps) @@ -1476,6 +1478,36 @@ def test_video_backend_handles_edit_list_trimmed_video( assert qwen_frames.shape[0] >= 4 +def test_opencv_demux_counts_frames_when_video_header_lacks_metadata( + monkeypatch: pytest.MonkeyPatch, +): + """Streamed WebM and raw H.264 carry no frame count/duration in their + headers, so OpenCV reports a bogus negative ``CAP_PROP_FRAME_COUNT``. + The loader must count the frames via a demux-only pass, without + decoding. + """ + with monkeypatch.context() as m: + m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv") + loader = VIDEO_LOADER_REGISTRY.load("opencv") + + webm_data = create_live_webm_video(num_frames=20, fps=10) + frames, metadata = loader.load_bytes(webm_data, backend="opencv") + assert frames.shape[0] == 20 + assert metadata["total_num_frames"] == 20 + assert metadata["duration"] == pytest.approx(2.0) + + h264_data = create_raw_h264_video(num_frames=20, fps=25) + frames, metadata = loader.load_bytes(h264_data, backend="opencv") + assert frames.shape[0] == 20 + assert metadata["total_num_frames"] == 20 + assert metadata["duration"] == pytest.approx(0.8) + + # The Qwen samplers used to raise on the bogus negative frame count. + qwen_frames, qwen_metadata = Qwen2VLVideoBackend.load_bytes(webm_data) + assert qwen_metadata["total_num_frames"] == 20 + assert qwen_frames.shape[0] >= 4 + + # ============================================================================ # Frame Recovery Tests # ============================================================================ diff --git a/tests/multimodal/utils.py b/tests/multimodal/utils.py index 23a05a0944..d74f684531 100644 --- a/tests/multimodal/utils.py +++ b/tests/multimodal/utils.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import io + import cv2 import numpy as np import numpy.typing as npt @@ -77,8 +79,6 @@ def create_long_gop_video( The marker lets a test recover which frame the decoder actually returned, independent of any metadata label. """ - import io - import av buf = io.BytesIO() @@ -103,6 +103,63 @@ def create_long_gop_video( return buf.getvalue() +def create_live_webm_video( + num_frames: int = 20, + fps: int = 10, + width: int = 64, + height: int = 64, +) -> bytes: + """Encode a WebM stream without container duration or frame-count metadata. + + The ``live`` muxing option mimics writing to a pipe: the Segment size is + unknown and no Duration or Cues elements are written, so decoder backends + cannot learn the frame count from the header. + """ + import av + + buf = io.BytesIO() + with av.open(buf, mode="w", format="webm", options={"live": "1"}) as container: + stream = container.add_stream("libvpx", rate=fps) + stream.width = width + stream.height = height + stream.pix_fmt = "yuv420p" + for i in range(num_frames): + img = np.zeros((height, width, 3), dtype=np.uint8) + img[:, :, 1] = i % 256 + frame = av.VideoFrame.from_ndarray(img, format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + return buf.getvalue() + + +def create_raw_h264_video( + num_frames: int = 20, + fps: int = 25, + width: int = 64, + height: int = 64, +) -> bytes: + """Encode a raw H.264 Annex B elementary stream (no container metadata).""" + import av + + buf = io.BytesIO() + with av.open(buf, mode="w", format="h264") as container: + stream = container.add_stream("h264", rate=fps) + stream.width = width + stream.height = height + stream.pix_fmt = "yuv420p" + for i in range(num_frames): + img = np.zeros((height, width, 3), dtype=np.uint8) + img[:, :, 1] = i % 256 + frame = av.VideoFrame.from_ndarray(img, format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + return buf.getvalue() + + def create_edit_list_trimmed_video( num_frames: int = 90, trim_start_frame: int = 60, @@ -117,8 +174,6 @@ def create_edit_list_trimmed_video( trailing frames are presentable — the shape produced by lossless trims (e.g. ``ffmpeg -ss ... -c copy``). Returns (video_bytes, visible_frames). """ - import io - import av src = create_long_gop_video(num_frames=num_frames, fps=fps) diff --git a/tests/samplers/test_beam_search.py b/tests/samplers/test_beam_search.py index 823202ba32..4e00c41559 100644 --- a/tests/samplers/test_beam_search.py +++ b/tests/samplers/test_beam_search.py @@ -99,6 +99,7 @@ def test_beam_search_abort_returns_partial_outputs_and_continues_other_prompts( llm = LLM.__new__(LLM) llm.llm_engine = Mock() + llm.llm_engine.input_processor.generation_config_fields = {} tokenizer = SimpleNamespace( eos_token_id=0, decode=lambda tokens, skip_special_tokens=False: " ".join( diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py index a6e4967620..b8134e6c8d 100644 --- a/tests/samplers/test_beam_search_online.py +++ b/tests/samplers/test_beam_search_online.py @@ -67,6 +67,7 @@ class _VllmConfig: class _InputProcessor: vllm_config = _VllmConfig() + generation_config_fields: dict[str, Any] = {} def resolve_watermarking(self, params): return self.vllm_config._check_supports_watermarking(params) @@ -128,7 +129,7 @@ class _OfflineServing(BeamSearchOfflineMixin): self, instances_batch: list[BeamSearchInstance], base_sampling_params: SamplingParams, - eos_token_id: int | None, + eos_token_ids: set[int], ignore_eos: bool, beam_width: int, sort_beams_key: Callable[..., Any], @@ -140,6 +141,74 @@ class _OfflineServing(BeamSearchOfflineMixin): return True +@pytest.mark.asyncio +@pytest.mark.parametrize("offline", [False, True]) +@pytest.mark.parametrize("config_eos", [None, 11, [0, 11]]) +@pytest.mark.parametrize("ignore_eos", [False, True]) +async def test_beam_search_respects_generation_config_eos( + monkeypatch, offline: bool, config_eos: int | list[int] | None, ignore_eos: bool +) -> None: + """Model-specific end-of-turn tokens must finish beams like tokenizer EOS.""" + monkeypatch.setattr( + _InputProcessor, "generation_config_fields", {"eos_token_id": config_eos} + ) + prompt: TokensInput = { + "type": "token", + "prompt_token_ids": [1], + } + params = BeamSearchParams( + beam_width=1, + max_tokens=3, + ignore_eos=ignore_eos, + include_stop_str_in_output=True, + watermarking=False, + ) + result = await anext( + _EngineClient().generate(prompt, SamplingParams(watermarking=False)) + ) + result.outputs[0].logprobs = [{11: Logprob(-0.1)}] + should_stop = config_eos is not None and not ignore_eos + expected_tokens = [11] if should_stop else [11, 11, 11] + engine_params: list[SamplingParams] = [] + + if offline: + offline_serving = _OfflineServing() + monkeypatch.setattr( + _OfflineServing, + "_beam_search_step", + BeamSearchOfflineMixin._beam_search_step, + ) + + def render_and_run_requests(*, params, **kwargs): + engine_params.extend(params) + return [result] + + monkeypatch.setattr( + offline_serving, "_render_and_run_requests", render_and_run_requests + ) + beam = offline_serving.beam_search(cast(list, [prompt]), params)[0].sequences[0] + assert beam.tokens == [1, *expected_tokens] + else: + + async def generate(engine_prompt, sampling_params, *args, **kwargs): + engine_params.append(sampling_params) + yield result + + online_serving = _AsyncServing() + monkeypatch.setattr(online_serving.engine_client, "generate", generate) + output = ( + await anext(online_serving.beam_search(prompt, "request", params)) + ).outputs[0] + assert output.token_ids == expected_tokens + assert output.finish_reason == ("stop" if should_stop else "length") + assert output.stop_reason == (11 if should_stop else None) + + # The engine validates stop_token_ids against the vocab before it applies + # the generation config, so the config EOS ids must not be copied into them. + assert engine_params + assert all(not p.stop_token_ids for p in engine_params) + + @pytest.mark.asyncio async def test_beam_search_handles_extra_logprob_candidates() -> None: prompt: TokensInput = { diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 740043bfa0..e5812e6673 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -59,6 +59,7 @@ CHECK_IMPORTS = { "vllm/distributed/device_communicators/shm_tensor_arena.py", "vllm/distributed/weight_transfer/ipc_engine.py", "vllm/distributed/weight_transfer/clients.py", + "examples/rl/rlhf_http_ipc_no_vllm.py", "vllm/model_executor/model_loader/weight_cache/protocol.py", "tests/distributed/test_shm_broadcast.py", "tests/distributed/test_weight_transfer.py", diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index 441475cafd..53297a4736 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -27,6 +27,7 @@ from .utils import ( BeamSearchOutput, BeamSearchSequence, create_sort_beams_key_function, + get_beam_search_eos_token_ids, ) logger = init_logger(__name__) @@ -86,8 +87,6 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): self.llm_engine.input_processor.resolve_watermarking(params) tokenizer = self.renderer.get_tokenizer() - eos_token_id = tokenizer.eos_token_id - sort_beams_key = create_sort_beams_key_function(eos_token_id, length_penalty) engine_inputs = self._preprocess_cmpl(prompts) lora_requests = self._lora_request_to_seq(lora_request, len(engine_inputs)) @@ -129,6 +128,11 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): params.structured_outputs, tokenizer ) + eos_token_ids = get_beam_search_eos_token_ids( + self.llm_engine.input_processor.generation_config_fields, + tokenizer.eos_token_id, + ) + sort_beams_key = create_sort_beams_key_function(eos_token_ids, length_penalty) instances: list[BeamSearchInstance] = [] for lora_req, prompt in zip(lora_requests, engine_inputs): @@ -169,7 +173,7 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): should_stop = self._beam_search_step( instances_batch=instances_batch, base_sampling_params=base_sampling_params, - eos_token_id=eos_token_id, + eos_token_ids=eos_token_ids, ignore_eos=ignore_eos, beam_width=beam_width, sort_beams_key=sort_beams_key, @@ -204,7 +208,7 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): self, instances_batch: list[BeamSearchInstance], base_sampling_params: SamplingParams, - eos_token_id: int | None, + eos_token_ids: set[int], ignore_eos: bool, beam_width: int, sort_beams_key: Callable, @@ -350,7 +354,7 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): cum_logprob=current_beam.cum_logprob + logprob_obj.logprob, ) - if token_id == eos_token_id and not ignore_eos: + if token_id in eos_token_ids and not ignore_eos: instance.completed.append(new_beam) else: instance_new_beams.append(new_beam) diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index 9cbe4ea258..1df1be0145 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -20,6 +20,7 @@ from vllm.utils.async_utils import collect_from_async_generator from .utils import ( BeamSearchSequence, create_sort_beams_key_function, + get_beam_search_eos_token_ids, ) @@ -47,8 +48,6 @@ class BeamSearchOnlineMixin(ABC): self.engine_client.input_processor.resolve_watermarking(params) tokenizer = self.renderer.get_tokenizer() - eos_token_id = tokenizer.eos_token_id - sort_beams_key = create_sort_beams_key_function(eos_token_id, length_penalty) if prompt["type"] == "embeds": raise NotImplementedError("Embedding prompt not supported for beam search") @@ -74,6 +73,11 @@ class BeamSearchOnlineMixin(ABC): sampling_params, self.engine_client.input_processor.vllm_config ) allowed = set(params.allowed_token_ids) if params.allowed_token_ids else None + eos_token_ids = get_beam_search_eos_token_ids( + self.engine_client.input_processor.generation_config_fields, + tokenizer.eos_token_id, + ) + sort_beams_key = create_sort_beams_key_function(eos_token_ids, length_penalty) all_beams = [ BeamSearchSequence( orig_prompt=prompt, @@ -159,17 +163,17 @@ class BeamSearchOnlineMixin(ABC): candidate_logprob = ( current_beam.cum_logprob + logprob_obj.logprob ) - if token_id == eos_token_id and not ignore_eos: + if token_id in eos_token_ids and not ignore_eos: completed.append( BeamSearchSequence( orig_prompt=prompt, - tokens=current_beam.tokens + [eos_token_id] + tokens=current_beam.tokens + [token_id] if include_stop_str_in_output else current_beam.tokens, logprobs=current_beam.logprobs + [logprobs], cum_logprob=candidate_logprob, finish_reason="stop", - stop_reason=eos_token_id, + stop_reason=token_id, ) ) else: @@ -219,7 +223,7 @@ class BeamSearchOnlineMixin(ABC): best_beams = sorted_completed[:beam_width] for beam in best_beams: - if beam.tokens[-1] == eos_token_id and not ignore_eos: + if beam.tokens[-1] in eos_token_ids and not ignore_eos: # Skip the eos token in the text. tokens = beam.tokens[tokenized_length:-1] else: diff --git a/vllm/entrypoints/generate/beam_search/utils.py b/vllm/entrypoints/generate/beam_search/utils.py index 61f97963b6..b0312b75f4 100644 --- a/vllm/entrypoints/generate/beam_search/utils.py +++ b/vllm/entrypoints/generate/beam_search/utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass +from typing import Any from vllm.inputs import ( DecoderOnlyEngineInput, @@ -13,6 +14,7 @@ from vllm.inputs import ( ) from vllm.logprobs import Logprob from vllm.lora.request import LoRARequest +from vllm.sampling_params import SamplingParams @dataclass @@ -137,7 +139,7 @@ class BeamSearchInstance: def get_beam_search_score( tokens: list[int], cumulative_logprob: float, - eos_token_id: int, + eos_token_ids: set[int], length_penalty: float = 1.0, ) -> float: """Calculate the beam search score with length penalty. @@ -147,17 +149,31 @@ def get_beam_search_score( https://github.com/huggingface/transformers/blob/ccb92be23def445f2afdea94c31286f84b89eb5b/src/transformers/generation/beam_search.py#L938 """ seq_len = len(tokens) - if tokens[-1] == eos_token_id: + if tokens[-1] in eos_token_ids: seq_len -= 1 # An aborted beam may contain only an EOS prompt token. return cumulative_logprob / (max(seq_len, 1) ** length_penalty) -def create_sort_beams_key_function(eos_token_id: int, length_penalty: float): +def create_sort_beams_key_function(eos_token_ids: set[int], length_penalty: float): def sort_beams_key(x: BeamSearchSequence) -> float: return get_beam_search_score( - x.tokens, x.cum_logprob, eos_token_id, length_penalty + x.tokens, x.cum_logprob, eos_token_ids, length_penalty ) return sort_beams_key + + +def get_beam_search_eos_token_ids( + generation_config: dict[str, Any], eos_token_id: int | None +) -> set[int]: + """Return the token ids that end generation: the tokenizer EOS plus the + generation config's ``eos_token_id``. + + This uses a separate SamplingParams so the params sent to the engine keep + their own ``stop_token_ids`` and are validated as before. + """ + params = SamplingParams() + params.update_from_generation_config(generation_config, eos_token_id) + return params.all_stop_token_ids diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 7c4a1fe917..2613a4073a 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -9,14 +9,17 @@ import torch.nn.functional as F from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt from vllm.exceptions import ChatTemplateResolutionError from vllm.logger import init_logger -from vllm.renderers import TokenizeParams -from vllm.renderers.hf import safe_apply_chat_template +from vllm.renderers import TokenizeParams, merge_kwargs +from vllm.renderers.hf import resolve_chat_template, safe_apply_chat_template from vllm.renderers.inputs.preprocess import ( extract_target_prompt, parse_model_prompt, prompt_to_seq, ) from vllm.tasks import PoolingTask +from vllm.transformers_utils.config import ( + get_sentence_transformers_chat_template_kwargs, +) from vllm.utils.mistral import is_mistral_tokenizer from ..base.io_processor import PoolingIOProcessor @@ -457,9 +460,22 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): self.model = model if self.supports_score_template else None self.use_sep_token = self.model_config.use_sep_token + self.saved_chat_template: str | None = None + self.saved_chat_template_kwargs: dict[str, Any] = {} + if not self.supports_score_template: + try: + self.saved_chat_template, self.saved_chat_template_kwargs = ( + self._get_saved_chat_template() + ) + except ValueError: + # An explicit chat template takes precedence anyway. + if self.chat_template is None: + raise + if ( getattr(self.model_config.hf_config, "is_original_qwen3_reranker", False) and self.chat_template is None + and self.saved_chat_template is None ): suggested_template = ( "examples/pooling/score/template/qwen3_vl_reranker.jinja" @@ -475,6 +491,41 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): suggested_template, ) + def _get_saved_chat_template(self) -> tuple[str | None, dict[str, Any]]: + """Get the chat template and kwargs declared by the Sentence + Transformers config, if any.""" + saved_kwargs = get_sentence_transformers_chat_template_kwargs( + self.model_config.model, self.model_config.revision + ) + if saved_kwargs is None: + return None, {} + + saved_kwargs = dict(saved_kwargs) + name = saved_kwargs.pop("chat_template", None) + if name == "default": + name = None + template = resolve_chat_template( + self.tokenizer, + chat_template=name, + tools=None, + model_config=self.model_config, + ) + # Hugging Face returns an unknown template name unchanged. + if template is None or template == name: + raise ValueError( + "The Sentence Transformers config declares message inputs, but " + f"the checkpoint has no {name or 'default'!r} chat template." + ) + + if self.chat_template is None: + logger.info_once( + "Scoring with the %r chat template declared by the Sentence " + "Transformers config. Its suffix is not restored after " + "truncation; use max_tokens_per_doc to truncate documents.", + name or "default", + ) + return template, saved_kwargs + ####################################### # online APIs @@ -648,6 +699,12 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): model_config = self.model_config tokenizer = self.tokenizer + if chat_template is None and self.saved_chat_template is not None: + chat_template = self.saved_chat_template + chat_template_kwargs = merge_kwargs( + self.saved_chat_template_kwargs, chat_template_kwargs + ) + prompt_1, prompt_2, mm_data, mm_uuids = parse_score_data( data_1, data_2, @@ -703,7 +760,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): prompt_inputs = tokenizer(text=full_prompt, **local_kwargs) return full_prompt, prompt_inputs - # FIXME: For now, we only apply a template when one is explicitly provided. + # FIXME: For now, we only apply a template when one is explicitly provided, + # either directly or by the Sentence Transformers config. # We cannot rely on the tokenizer's chat template because many models # inherit junk templates from their base LLM, which breaks both the models # and the tests that use them. diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index d1a8ca1d9d..569cc0c9f5 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -82,6 +82,8 @@ from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950, on_gfx1250 from vllm.sequence import IntermediateTensors +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel +from vllm.utils.torch_utils import aux_stream from vllm.v1.worker.ubatching import dbo_current_ubatch_token_offset from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID @@ -713,6 +715,18 @@ class DeepseekV4MoE(nn.Module): ), prefix=f"{prefix}.experts", ) + self.shared_experts_stream = ( + aux_stream() + if self.shared_experts is not None + and not self.experts.fuse_shared_expert + and not envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM + else None + ) + self.shared_experts_events = ( + (torch.cuda.Event(), torch.cuda.Event()) + if self.shared_experts_stream is not None + else None + ) return fuse_shared_into_routed = ( @@ -777,6 +791,30 @@ class DeepseekV4MoE(nn.Module): def _forward_mega_moe( self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None + ) -> torch.Tensor: + shared_experts = self.shared_experts + if shared_experts is None: + return self._forward_mega_moe_routed(hidden_states, input_ids) + events = self.shared_experts_events + if ( + events is None + or hidden_states.shape[0] > envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD + ): + return self._forward_mega_moe_routed( + hidden_states, input_ids + ) + shared_experts(hidden_states) + # The unfused shared expert overlaps router + MegaMoE in small steps. + routed, shared = maybe_execute_in_parallel( + fn0=lambda: self._forward_mega_moe_routed(hidden_states, input_ids), + fn1=lambda: shared_experts(hidden_states), + event0=events[0], + event1=events[1], + aux_stream=self.shared_experts_stream, + ) + return routed + shared + + def _forward_mega_moe_routed( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None ) -> torch.Tensor: router_logits, _ = self.gate(hidden_states) topk_weights, topk_ids = fused_topk_bias( @@ -792,12 +830,7 @@ class DeepseekV4MoE(nn.Module): bias_vl=self.gate.bias_vl, image_sentinel_lo=self.image_sentinel_lo, ) - final_hidden_states = self.experts(hidden_states, topk_weights, topk_ids) - if self.shared_experts is not None: - final_hidden_states = final_hidden_states + self.shared_experts( - hidden_states - ) - return final_hidden_states + return self.experts(hidden_states, topk_weights, topk_ids) # Hidden sizes supported by AITER mhc_pre_big_fuse_rmsnorm. diff --git a/vllm/models/qwen4_exp/amd/hyperconnection.py b/vllm/models/qwen4_exp/amd/hyperconnection.py index 087f03fdcf..b6376ac238 100644 --- a/vllm/models/qwen4_exp/amd/hyperconnection.py +++ b/vllm/models/qwen4_exp/amd/hyperconnection.py @@ -70,7 +70,7 @@ class GatedResidual(nn.Module): ) -> None: super().__init__() self.config = config - self.lora_rank = config.hc_lowrank + self.hc_lowrank = config.hc_lowrank self.hc_count = config.hc_count self.hidden_size = config.hidden_size self.use_combine = use_combine @@ -91,11 +91,11 @@ class GatedResidual(nn.Module): # -- vLLM Linear weights -------------------------------------------- # The merged skinny-GEMM shape is physically padded to 16 rows for # alignment and efficient backend dispatch. - self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0 + self.pad_size = (-(self.hc_lowrank + self.hc_count)) % 16 if use_combine else 0 if use_combine: self.input_mix_weight_down_block_inject = MergedColumnParallelLinear( self.hyper_hidden_size, - [self.lora_rank, self.hc_count] + [self.hc_lowrank, self.hc_count] + ([self.pad_size] if self.pad_size else []), bias=False, params_dtype=config.params_dtype, @@ -107,7 +107,7 @@ class GatedResidual(nn.Module): else: self.input_mix_weight_down = ReplicatedLinear( self.hyper_hidden_size, - self.lora_rank, + self.hc_lowrank, bias=False, params_dtype=config.params_dtype, quant_config=None, @@ -115,7 +115,7 @@ class GatedResidual(nn.Module): return_bias=False, ) self.input_mix_weight_up = ReplicatedLinear( - self.lora_rank, + self.hc_lowrank, self.hyper_hidden_size, bias=False, params_dtype=config.params_dtype, @@ -136,15 +136,15 @@ class GatedResidual(nn.Module): if self.use_combine: # produce injection logits for combine - split_sizes = [self.lora_rank, self.hc_count, self.pad_size] + split_sizes = [self.hc_lowrank, self.hc_count, self.pad_size] down_and_injection = self.input_mix_weight_down_block_inject(xn) - lora, injection, _ = down_and_injection.split(split_sizes, dim=-1) + lowrank, injection, _ = down_and_injection.split(split_sizes, dim=-1) else: - lora = self.input_mix_weight_down(xn) + lowrank = self.input_mix_weight_down(xn) injection = None - lora = hc_silu(lora, self.hc_count) - gate = self.input_mix_weight_up(lora) # [M, D] + lowrank = hc_silu(lowrank, self.hc_count) + gate = self.input_mix_weight_up(lowrank) # [M, D] block_input = hc_gate_mix(xn, gate, self.hc_count) return hidden_states, block_input, injection @@ -172,15 +172,15 @@ class GatedResidual(nn.Module): if self.use_combine: # produce injection logits for combine - split_sizes = [self.lora_rank, self.hc_count, self.pad_size] + split_sizes = [self.hc_lowrank, self.hc_count, self.pad_size] down_and_injection = self.input_mix_weight_down_block_inject(xn) - lora, injection, _ = down_and_injection.split(split_sizes, dim=-1) + lowrank, injection, _ = down_and_injection.split(split_sizes, dim=-1) else: - lora = self.input_mix_weight_down(xn) + lowrank = self.input_mix_weight_down(xn) injection = None - lora = hc_silu(lora, self.hc_count) - gate = self.input_mix_weight_up(lora) # [M, D] + lowrank = hc_silu(lowrank, self.hc_count) + gate = self.input_mix_weight_up(lowrank) # [M, D] block_input = hc_gate_mix(xn, gate, self.hc_count) return hidden_states, block_input, injection diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py index c1cfc0f246..478b0ae972 100644 --- a/vllm/models/qwen4_exp/amd/ple_layer.py +++ b/vllm/models/qwen4_exp/amd/ple_layer.py @@ -29,6 +29,7 @@ from vllm.models.common.ops.sequence_parallel import ( sp_reduce_scatter, sp_shard, ) +from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.short_conv_attn import ( @@ -242,7 +243,7 @@ class Qwen4ExpNGramEmbedding(nn.Module): persistent=True, ) divisor = int(config.make_ngram_vocab_size_divisible_by) - padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor + padded_vocab_size = round_up(total_vocab_size, divisor) if params_dtype is None: params_dtype = torch.get_default_dtype() embedding_prefix = f"{prefix}.ngram_embedding" @@ -280,9 +281,11 @@ class Qwen4ExpNGramEmbedding(nn.Module): reduce_results=not self.use_reduce_scatter, ) logger.info( - "Initialized AMD PLE embedding %s: quantization_method=%s, " + "Initialized AMD PLE embedding %s: " + "quantization_config=%s, quantization_method=%s, " "weight_dtype=%s, weight_device=%s, pinned=%s", embedding_prefix, + type(quant_config).__name__, type(embedding_quant_method).__name__, self.ngram_embedding.weight.dtype, self.ngram_embedding.weight.device, @@ -742,7 +745,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): ) kv, _ = self.kv_proj(embeddings) key, value = kv.split(self.kv_proj.output_sizes, dim=-1) - gated_output, conv_input = ple_gate( + gated_output, conv_input = torch.ops.vllm.qwen4_exp_ple_gate( key, value, hidden_states, @@ -848,6 +851,28 @@ def qwen4_exp_ple_short_conv( layer._short_conv(inputs, residual, outer_residual) +def _ple_gate_fake( + key: torch.Tensor, + value: torch.Tensor, + hidden: torch.Tensor, + norm_key_w: torch.Tensor, + norm_query_w: torch.Tensor, + norm_conv_w: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return the gate output shapes for AMD graph tracing.""" + return torch.empty_like(hidden), torch.empty_like(hidden) + + +# Keep the AMD gate opaque to Dynamo while reusing the shared Triton kernel. +direct_register_custom_op( + op_name="qwen4_exp_ple_gate", + op_func=ple_gate, + mutates_args=[], + fake_impl=_ple_gate_fake, +) + + direct_register_custom_op( op_name="qwen4_exp_amd_ple_ngram_ids", op_func=qwen4_exp_amd_ple_ngram_ids, diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py index 0d875d015c..8e9a04ceed 100644 --- a/vllm/models/qwen4_exp/common/ngram_embedding.py +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -194,22 +194,21 @@ class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): ) if quant_config is None: return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if isinstance(quant_config, ModelOptMixedPrecisionConfig): - quant_algo = quant_config._resolve_quant_algo(prefix) - if quant_algo == "FP8": - return Qwen4ExpPLEFp8EmbeddingMethod() - if quant_algo == "NVFP4": - return Qwen4ExpPLENvFp4EmbeddingMethod() - if quant_algo is not None: - raise ValueError( - f"Unsupported quant_algo {quant_algo!r} for the PLE embedding " - f"{prefix!r}; expected FP8 or NVFP4" - ) - return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if isinstance( - quant_config, ModelOptQuantConfigBase - ) and quant_config.is_layer_excluded(prefix): - return Qwen4ExpPLEUnquantizedEmbeddingMethod() + if isinstance(quant_config, ModelOptQuantConfigBase): + if quant_config.is_layer_excluded(prefix): + return Qwen4ExpPLEUnquantizedEmbeddingMethod() + if isinstance(quant_config, ModelOptMixedPrecisionConfig): + quant_algo = quant_config._resolve_quant_algo(prefix) + if quant_algo == "FP8": + return Qwen4ExpPLEFp8EmbeddingMethod() + if quant_algo == "NVFP4": + return Qwen4ExpPLENvFp4EmbeddingMethod() + if quant_algo is not None: + raise ValueError( + f"Unsupported quant_algo {quant_algo!r} for the PLE embedding " + f"{prefix!r}; expected FP8 or NVFP4" + ) + return Qwen4ExpPLEUnquantizedEmbeddingMethod() if ( isinstance(quant_config, CompressedTensorsConfig) and quant_config.get_scheme_dict(None, layer_name=prefix) is None @@ -223,29 +222,32 @@ class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): # Quark quantizes only Linear and MoE layers; PLE tables stay BF16. if isinstance(quant_config, QuarkConfig): return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if not isinstance(quant_config, Fp8Config): - raise NotImplementedError( - "Qwen4Exp PLE embedding does not support quantization config " - f"{type(quant_config).__name__}" - ) - - ignored_layers = quant_config.ignored_layers - if is_layer_skipped( - prefix, - ignored_layers, - quant_config.packed_modules_mapping, - match_mode=quant_config.ignored_layers_match_mode, - ): - return Qwen4ExpPLEUnquantizedEmbeddingMethod() - # PLE checkpoint shards form one runtime embedding parameter. - shard_prefix = f"{prefix}.shard_" - if any(name.startswith(shard_prefix) for name in ignored_layers): - return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if not quant_config.is_checkpoint_fp8_serialized: + if isinstance(quant_config, Fp8Config): + ignored_layers = quant_config.ignored_layers + if is_layer_skipped( + prefix, + ignored_layers, + quant_config.packed_modules_mapping, + match_mode=quant_config.ignored_layers_match_mode, + ): + return Qwen4ExpPLEUnquantizedEmbeddingMethod() + # PLE checkpoint shards form one runtime embedding parameter. + shard_prefix = f"{prefix}.shard_" + if any(name.startswith(shard_prefix) for name in ignored_layers): + return Qwen4ExpPLEUnquantizedEmbeddingMethod() + if quant_config.is_checkpoint_fp8_serialized: + return Qwen4ExpPLEFp8EmbeddingMethod() raise NotImplementedError( "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints" ) - return Qwen4ExpPLEFp8EmbeddingMethod() + + logger.warning( + "Falling back to unquantized PLE embedding for %s " + "with quantization config %s.", + prefix, + type(quant_config).__name__, + ) + return Qwen4ExpPLEUnquantizedEmbeddingMethod() def apply( self, @@ -781,6 +783,7 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): ) active_output = buffer[: gathered_ids.shape[0]] prefetch_stream.wait_stream(torch.cuda.current_stream()) + # Prevent ID storage reuse until the prefetch stream finishes reading it. gathered_ids.record_stream(prefetch_stream) with torch.cuda.stream(prefetch_stream): self._lookup(gathered_ids, output=active_output) diff --git a/vllm/models/qwen4_exp/common/ops/ple.py b/vllm/models/qwen4_exp/common/ops/ple.py index 41ce082c9e..994b799ecd 100644 --- a/vllm/models/qwen4_exp/common/ops/ple.py +++ b/vllm/models/qwen4_exp/common/ops/ple.py @@ -16,7 +16,6 @@ import torch from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @@ -240,7 +239,7 @@ def _ple_gate_kernel( tl.store(normed_ptr + t * HC * H + offs, normed, mask=mask) -def _ple_gate( +def ple_gate( key: torch.Tensor, value: torch.Tensor, hidden: torch.Tensor, @@ -281,40 +280,6 @@ def _ple_gate( return gated, normed -def _ple_gate_fake( - key: torch.Tensor, - value: torch.Tensor, - hidden: torch.Tensor, - norm_key_w: torch.Tensor, - norm_query_w: torch.Tensor, - norm_conv_w: torch.Tensor, - eps: float, -) -> tuple[torch.Tensor, torch.Tensor]: - return torch.empty_like(hidden), torch.empty_like(hidden) - - -direct_register_custom_op( - op_name="qwen4_exp_ple_gate", - op_func=_ple_gate, - mutates_args=[], - fake_impl=_ple_gate_fake, -) - - -def ple_gate( - key: torch.Tensor, - value: torch.Tensor, - hidden: torch.Tensor, - norm_key_w: torch.Tensor, - norm_query_w: torch.Tensor, - norm_conv_w: torch.Tensor, - eps: float, -) -> tuple[torch.Tensor, torch.Tensor]: - return torch.ops.vllm.qwen4_exp_ple_gate( - key, value, hidden, norm_key_w, norm_query_w, norm_conv_w, eps - ) - - # Dilated short convolution @triton.jit(do_not_specialize=["num_reqs", "bs_iters", "has_token_map"]) def _ple_conv_kernel( diff --git a/vllm/models/qwen4_exp/common/ple.py b/vllm/models/qwen4_exp/common/ple.py index 72fc494b60..e15e803b69 100644 --- a/vllm/models/qwen4_exp/common/ple.py +++ b/vllm/models/qwen4_exp/common/ple.py @@ -20,66 +20,64 @@ class PLEShardOverlap: row_count: int -def compute_ple_shard_overlap( - *, - checkpoint_start: int, - checkpoint_rows: int, - tp_start: int, - tp_end: int, -) -> PLEShardOverlap | None: - """Compute the overlap of a checkpoint shard and one TP vocabulary range.""" - if checkpoint_start < 0 or checkpoint_rows < 0: - raise ValueError("checkpoint shard bounds must be non-negative") - if tp_start < 0 or tp_end < tp_start: - raise ValueError("invalid TP vocabulary range") - checkpoint_end = checkpoint_start + checkpoint_rows - overlap_start = max(checkpoint_start, tp_start) - overlap_end = min(checkpoint_end, tp_end) - if overlap_start >= overlap_end: - return None - return PLEShardOverlap( - source_start=overlap_start - checkpoint_start, - destination_start=overlap_start - tp_start, - row_count=overlap_end - overlap_start, - ) - - -def copy_ple_embedding_shard_( - destination: torch.Tensor, - loaded_weight: torch.Tensor, - *, - checkpoint_start: int, - tp_start: int, - tp_end: int, -) -> int: - """Copy the overlapping rows of a PLE checkpoint shard into a TP table.""" - if destination.ndim == 0 or loaded_weight.ndim != destination.ndim: - raise ValueError("destination and loaded weight must have matching ranks") - if destination.shape[1:] != loaded_weight.shape[1:]: - raise ValueError( - "embedding shard dimensions do not match: " - f"{tuple(destination.shape[1:])} != {tuple(loaded_weight.shape[1:])}" - ) - if destination.shape[0] < tp_end - tp_start: - raise ValueError("destination does not cover the requested TP range") - overlap = compute_ple_shard_overlap( - checkpoint_start=checkpoint_start, - checkpoint_rows=loaded_weight.shape[0], - tp_start=tp_start, - tp_end=tp_end, - ) - if overlap is None: - return 0 - source = loaded_weight.narrow(0, overlap.source_start, overlap.row_count) - target = destination.narrow(0, overlap.destination_start, overlap.row_count) - with torch.no_grad(): - target.copy_(source.to(device=target.device, dtype=target.dtype)) - return overlap.row_count - - class PLEVocabParallelEmbedding(VocabParallelEmbedding): """Vocab-parallel embedding that accepts checkpoint row shards.""" + def _compute_shard_overlap( + self, + *, + checkpoint_start: int, + checkpoint_rows: int, + ) -> PLEShardOverlap | None: + """Compute the overlap of a checkpoint shard and one TP vocabulary range.""" + tp_start = self.shard_indices.org_vocab_start_index + tp_end = self.shard_indices.org_vocab_end_index + if checkpoint_start < 0 or checkpoint_rows < 0: + raise ValueError("checkpoint shard bounds must be non-negative") + if tp_start < 0 or tp_end < tp_start: + raise ValueError("invalid TP vocabulary range") + checkpoint_end = checkpoint_start + checkpoint_rows + overlap_start = max(checkpoint_start, tp_start) + overlap_end = min(checkpoint_end, tp_end) + if overlap_start >= overlap_end: + return None + return PLEShardOverlap( + source_start=overlap_start - checkpoint_start, + destination_start=overlap_start - tp_start, + row_count=overlap_end - overlap_start, + ) + + def _copy_embedding_shard_( + self, + destination: torch.Tensor, + loaded_weight: torch.Tensor, + *, + checkpoint_start: int, + ) -> int: + """Copy the overlapping rows of a PLE checkpoint shard into a TP table.""" + tp_start = self.shard_indices.org_vocab_start_index + tp_end = self.shard_indices.org_vocab_end_index + if destination.ndim == 0 or loaded_weight.ndim != destination.ndim: + raise ValueError("destination and loaded weight must have matching ranks") + if destination.shape[1:] != loaded_weight.shape[1:]: + raise ValueError( + "embedding shard dimensions do not match: " + f"{tuple(destination.shape[1:])} != {tuple(loaded_weight.shape[1:])}" + ) + if destination.shape[0] < tp_end - tp_start: + raise ValueError("destination does not cover the requested TP range") + overlap = self._compute_shard_overlap( + checkpoint_start=checkpoint_start, + checkpoint_rows=loaded_weight.shape[0], + ) + if overlap is None: + return 0 + source = loaded_weight.narrow(0, overlap.source_start, overlap.row_count) + target = destination.narrow(0, overlap.destination_start, overlap.row_count) + with torch.no_grad(): + target.copy_(source.to(device=target.device, dtype=target.dtype)) + return overlap.row_count + def weight_loader( self, param: torch.Tensor, @@ -89,10 +87,8 @@ class PLEVocabParallelEmbedding(VocabParallelEmbedding): if checkpoint_start is None: super().weight_loader(param, loaded_weight) return - copy_ple_embedding_shard_( + self._copy_embedding_shard_( param, loaded_weight, checkpoint_start=checkpoint_start, - tp_start=self.shard_indices.org_vocab_start_index, - tp_end=self.shard_indices.org_vocab_end_index, ) diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py index ea37eb5cc6..df088ad2fc 100644 --- a/vllm/models/qwen4_exp/common/qsa_cache.py +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -24,7 +24,7 @@ from vllm.config.cache import CacheDType from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton -from vllm.utils.math_utils import cdiv +from vllm.utils.math_utils import cdiv, round_up from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -880,9 +880,7 @@ class QSAKeyStateCache(_QSAStateCache): key_head_size = int(kwargs.pop("head_size")) self.key_head_size = key_head_size self.cache_rope_positions = bool(cache_rope_positions) - self.rope_position_offset = ( - (key_head_size + self._BF16_PER_INT64 - 1) // self._BF16_PER_INT64 - ) * self._BF16_PER_INT64 + self.rope_position_offset = round_up(key_head_size, self._BF16_PER_INT64) storage_head_size = key_head_size if self.cache_rope_positions: storage_head_size = self.rope_position_offset + ( @@ -910,7 +908,7 @@ class QSAKeyStateCache(_QSAStateCache): # the scheduler block size). Anything narrower lets a rejected draft row # overwrite a committed key the next step needs to close the group. span = self.compress_ratio + vllm_config.num_speculative_tokens - capacity = self.compress_ratio * cdiv(span, self.compress_ratio) + capacity = round_up(span, self.compress_ratio) assert self.cache_config.block_size % capacity == 0, ( f"QSA ring capacity {capacity} must divide the attention block " f"size {self.cache_config.block_size}" diff --git a/vllm/models/qwen4_exp/nvidia/hyperconnection.py b/vllm/models/qwen4_exp/nvidia/hyperconnection.py index dcf8daa956..3844f6de32 100644 --- a/vllm/models/qwen4_exp/nvidia/hyperconnection.py +++ b/vllm/models/qwen4_exp/nvidia/hyperconnection.py @@ -74,7 +74,7 @@ class GatedResidual(nn.Module): ) -> None: super().__init__() self.config = config - self.lora_rank = config.hc_lowrank + self.hc_lowrank = config.hc_lowrank self.hc_count = config.hc_count self.hidden_size = config.hidden_size self.use_combine = use_combine @@ -95,11 +95,11 @@ class GatedResidual(nn.Module): # -- vLLM Linear weights -------------------------------------------- # The merged skinny-GEMM shape is physically padded to 16 rows to ensure # good alignment and performant implementation chosen by CuBLAS heuristics. - self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0 + self.pad_size = (-(self.hc_lowrank + self.hc_count)) % 16 if use_combine else 0 if use_combine: self.input_mix_weight_down_block_inject = MergedColumnParallelLinear( self.hyper_hidden_size, - [self.lora_rank, self.hc_count] + [self.hc_lowrank, self.hc_count] + ([self.pad_size] if self.pad_size else []), bias=False, params_dtype=config.params_dtype, @@ -117,7 +117,7 @@ class GatedResidual(nn.Module): else: self.input_mix_weight_down = ReplicatedLinear( self.hyper_hidden_size, - self.lora_rank, + self.hc_lowrank, bias=False, params_dtype=config.params_dtype, quant_config=None, @@ -125,7 +125,7 @@ class GatedResidual(nn.Module): return_bias=False, ) self.input_mix_weight_up = ReplicatedLinear( - self.lora_rank, + self.hc_lowrank, self.hyper_hidden_size, bias=False, params_dtype=config.params_dtype, @@ -150,13 +150,13 @@ class GatedResidual(nn.Module): return hc_down_silu( xn, self.input_mix_weight_down_block_inject.weight, - self.lora_rank, + self.hc_lowrank, self.hc_count, ) down_and_injection = self.input_mix_weight_down_block_inject(xn) - split_sizes = [self.lora_rank, self.hc_count, self.pad_size] - lora, injection, _ = down_and_injection.split(split_sizes, dim=-1) - return hc_silu(lora, self.hc_count), injection + split_sizes = [self.hc_lowrank, self.hc_count, self.pad_size] + lowrank, injection, _ = down_and_injection.split(split_sizes, dim=-1) + return hc_silu(lowrank, self.hc_count), injection def mix( self, hidden_states: torch.Tensor @@ -168,8 +168,8 @@ class GatedResidual(nn.Module): self.hc_count, ) - lora, injection = self._down_and_inject(xn) - gate = self.input_mix_weight_up(lora) # [M, D] + lowrank, injection = self._down_and_inject(xn) + gate = self.input_mix_weight_up(lowrank) # [M, D] block_input = hc_gate_mix(xn, gate, self.hc_count) return hidden_states, block_input, injection @@ -196,8 +196,8 @@ class GatedResidual(nn.Module): self.hc_count, ) - lora, injection = self._down_and_inject(xn) - gate = self.input_mix_weight_up(lora) # [M, D] + lowrank, injection = self._down_and_inject(xn) + gate = self.input_mix_weight_up(lowrank) # [M, D] block_input = hc_gate_mix(xn, gate, self.hc_count) return hidden_states, block_input, injection diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index 0d1ae0e9b8..47c675a477 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.quantization.base_config import ( from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 from vllm.model_executor.models.utils import AutoWeightsLoader from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter, sp_shard +from vllm.utils.math_utils import round_up from ..common.ngram_embedding import ( Qwen4ExpPLEDeviceEmbedding, @@ -356,7 +357,7 @@ class Qwen4ExpNGramEmbedding(nn.Module): persistent=True, ) divisor = int(config.make_ngram_vocab_size_divisible_by) - padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor + padded_vocab_size = round_up(total_vocab_size, divisor) embedding_prefix = f"{prefix}.ngram_embedding" embedding_quant_method = Qwen4ExpPLEEmbeddingMethod.from_quant_config( quant_config, @@ -411,9 +412,11 @@ class Qwen4ExpNGramEmbedding(nn.Module): ) weight = self.ngram_embedding.weight logger.info( - "Initialized PLE embedding %s: quantization_method=%s, " + "Initialized PLE embedding %s: " + "quantization_config=%s, quantization_method=%s, " "weight_dtype=%s, weight_device=%s, pinned=%s", embedding_prefix, + type(quant_config).__name__, type(embedding_quant_method).__name__, weight.dtype, weight.device, diff --git a/vllm/models/qwen4_exp/nvidia/ops/hc.py b/vllm/models/qwen4_exp/nvidia/ops/hc.py index 36bc8cefbc..3dee9db436 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/hc.py +++ b/vllm/models/qwen4_exp/nvidia/ops/hc.py @@ -6,7 +6,6 @@ import torch from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.torch_utils import direct_register_custom_op @triton.jit @@ -52,7 +51,7 @@ def _grouped_gemma_rmsnorm_kernel( tl.store(y_ptr + row * stride_y + offsets, y, mask) -def _grouped_gemma_rmsnorm( +def grouped_gemma_rmsnorm( x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int ) -> torch.Tensor: N, DIM = x.shape @@ -105,7 +104,7 @@ def _hc_silu_kernel( tl.store(y_ptr + row * stride_y + offs, y, mask) -def _hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor: +def hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor: num_tokens, DIM = x.shape assert x.stride(1) == 1 @@ -160,7 +159,7 @@ def _hc_gate_mix_kernel( tl.store(y_ptr + row * stride_y + offs_inner, acc, mask) -def _hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor: +def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor: N, DIM = gate.shape assert x.shape == gate.shape assert DIM % hc_count == 0 @@ -232,7 +231,7 @@ def _hc_combine_kernel( tl.store(out_ptr + row * stride_out + offs, out, mask=mask) -def _hc_combine( +def hc_combine( residual: torch.Tensor, block_output: torch.Tensor, injection_logits: torch.Tensor | None, @@ -345,7 +344,7 @@ def _hc_combine_norm_kernel( tl.store(y_ptr + row * stride_y + offs, y, mask_inner) -def _hc_combine_norm( +def hc_combine_norm( residual: torch.Tensor, block_output: torch.Tensor, injection_logits: torch.Tensor | None, @@ -392,109 +391,6 @@ def _hc_combine_norm( return out, y -def _same_shape_fake(x: torch.Tensor, *args) -> torch.Tensor: - return x.new_empty(x.shape) - - -def _hc_gate_mix_fake( - x: torch.Tensor, gate: torch.Tensor, hc_count: int -) -> torch.Tensor: - del gate - return x.new_empty((x.shape[0], x.shape[1] // hc_count)) - - -def _hc_combine_fake( - residual: torch.Tensor, - block_output: torch.Tensor, - injection_logits: torch.Tensor | None, - hc_count: int, -) -> torch.Tensor: - del block_output, injection_logits, hc_count - return residual.new_empty(residual.shape) - - -def _hc_combine_norm_fake( - residual: torch.Tensor, - block_output: torch.Tensor, - injection_logits: torch.Tensor | None, - norm_weight: torch.Tensor, - eps: float, - hc_count: int, -) -> tuple[torch.Tensor, torch.Tensor]: - del block_output, injection_logits, norm_weight, eps, hc_count - return residual.new_empty(residual.shape), residual.new_empty(residual.shape) - - -direct_register_custom_op( - op_name="qwen4_exp_grouped_gemma_rmsnorm", - op_func=_grouped_gemma_rmsnorm, - fake_impl=_same_shape_fake, -) -direct_register_custom_op( - op_name="qwen4_exp_hc_silu", - op_func=_hc_silu, - fake_impl=_same_shape_fake, -) -direct_register_custom_op( - op_name="qwen4_exp_hc_gate_mix", - op_func=_hc_gate_mix, - fake_impl=_hc_gate_mix_fake, -) -direct_register_custom_op( - op_name="qwen4_exp_hc_combine", - op_func=_hc_combine, - fake_impl=_hc_combine_fake, -) -direct_register_custom_op( - op_name="qwen4_exp_hc_combine_norm", - op_func=_hc_combine_norm, - fake_impl=_hc_combine_norm_fake, -) - - -def grouped_gemma_rmsnorm( - x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int -) -> torch.Tensor: - return torch.ops.vllm.qwen4_exp_grouped_gemma_rmsnorm(x, weight, eps, num_groups) - - -def hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor: - return torch.ops.vllm.qwen4_exp_hc_silu(x, hc_count) - - -def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor: - return torch.ops.vllm.qwen4_exp_hc_gate_mix(x, gate, hc_count) - - -def hc_combine( - residual: torch.Tensor, - block_output: torch.Tensor, - injection_logits: torch.Tensor | None, - hc_count: int, -) -> torch.Tensor: - return torch.ops.vllm.qwen4_exp_hc_combine( - residual, block_output, injection_logits, hc_count - ) - - -def hc_combine_norm( - residual: torch.Tensor, - block_output: torch.Tensor, - injection_logits: torch.Tensor | None, - norm_weight: torch.Tensor, - eps: float, - hc_count: int, -) -> tuple[torch.Tensor, torch.Tensor]: - return torch.ops.vllm.qwen4_exp_hc_combine_norm( - residual, - block_output, - injection_logits, - norm_weight, - eps, - hc_count, - ) - - __all__ = [ "grouped_gemma_rmsnorm", "hc_combine", diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py index 9b28ed6dca..4aa85f79f3 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py @@ -10,6 +10,7 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import ( ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import round_up _TOPK_WORKSPACE_BYTES = 1024 * 1024 _DECODE_BLOCK_N = 64 @@ -611,7 +612,7 @@ def qsa_select_paged_prefill( # No row scores beyond cdiv(max_seq_len, compress_ratio) compressed # columns. Round up to 64 to keep the logits row stride # cooperative_topk-compatible. - logits_width = triton.cdiv(triton.cdiv(max_seq_len, compress_ratio), 64) * 64 + logits_width = round_up(triton.cdiv(max_seq_len, compress_ratio), 64) logits_width = min(max(64, logits_width), page_table.shape[1] * k_cache.shape[1]) # chunk the inputs to keep temp logits below VLLM_SPARSE_INDEXER_MAX_LOGITS_MB diff --git a/vllm/multimodal/video_decoders/opencv.py b/vllm/multimodal/video_decoders/opencv.py index 3ca141c738..88d7fb8375 100644 --- a/vllm/multimodal/video_decoders/opencv.py +++ b/vllm/multimodal/video_decoders/opencv.py @@ -38,7 +38,9 @@ def decode_opencv( int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), ) - source = loader_cls._prepare_source(OpenCVVideoBackendMixin.get_video_metadata(cap)) + source = loader_cls._prepare_source( + OpenCVVideoBackendMixin.get_video_metadata(data) + ) frame_idx = loader_cls.compute_frames_index_to_sample( source=source, target=target, **sampling_kwargs ) @@ -74,10 +76,34 @@ class OpenCVVideoBackendMixin: raise ValueError("Could not open video stream") return cap + @classmethod + def demux_count_frames(cls, data: bytes) -> int: + """Count frames with OpenCV's demuxer-only mode.""" + # NOTE: buf must stay referenced until after release(). + # OpenCV's Python stream binding does not safely own the stream object + buf = BytesIO(data) + demux_cap = cv2.VideoCapture( + buf, cls.get_cv2_video_api(), [cv2.CAP_PROP_FORMAT, -1] + ) + count = 0 + while demux_cap.grab(): + count += 1 + demux_cap.release() + if count == 0: + raise ValueError("Could not determine the number of video frames") + return count + @staticmethod - def get_video_metadata(cap: "cv2.VideoCapture") -> VideoSourceMetadata: + def get_video_metadata(data: bytes) -> VideoSourceMetadata: + cap = OpenCVVideoBackendMixin.open_video_capture(data) total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) original_fps = cap.get(cv2.CAP_PROP_FPS) + if total_frames_num <= 0: + # Streamed WebM/MKV and raw H.264 miss frame count in headers; + # We need to use demux frames count for retrieval + total_frames_num = OpenCVVideoBackendMixin.demux_count_frames(data) + duration = total_frames_num / original_fps if original_fps > 0 else 0 + return VideoSourceMetadata(total_frames_num, original_fps, duration) # CAP_PROP_FRAME_COUNT counts every physical sample in the container, # overstating the presentable frames when an mp4 edit list hides the # decode lead-in (e.g. a lossless `ffmpeg -ss ... -c copy` trim). @@ -87,7 +113,7 @@ class OpenCVVideoBackendMixin: # may land mid-stream on corrupt files, so only trust it when the # stream really ends there. A one-frame tolerance keeps mere header # miscounts from being "corrected". - if total_frames_num > 0 and cap.set(cv2.CAP_PROP_POS_AVI_RATIO, 1.0): + if cap.set(cv2.CAP_PROP_POS_AVI_RATIO, 1.0): visible_frames = int(cap.get(cv2.CAP_PROP_POS_FRAMES)) at_stream_end = not cap.grab() cap.set(cv2.CAP_PROP_POS_FRAMES, 0) diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index c67d468a13..8587c02b68 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -89,6 +89,10 @@ CONFIG_FIELD_ALIASES: dict[str, tuple[str, ...]] = { `ModelArchConfigConvertorBase` reads a scalar from these; a caller that needs the raw value (a per-layer `list`, say) passes the aliases to `getattr_iter` itself.""" +_ST_TRANSFORMER_MODULE_TYPES = { + "sentence_transformers.models.Transformer", + "sentence_transformers.base.modules.transformer.Transformer", +} _ST_POOLING_MODULE_TYPES = { "sentence_transformers.models.Pooling", "sentence_transformers.sentence_transformer.modules.pooling.Pooling", @@ -1354,6 +1358,88 @@ def try_get_dense_modules( return None +@cache +def get_sentence_transformers_chat_template_kwargs( + model: str | Path, + revision: str | None = None, +) -> dict[str, Any] | None: + """Get the saved chat template settings of a Sentence Transformers + CrossEncoder whose single Transformer module declares the `message` + modality. + + Sentence Transformers formats the pairs of such checkpoints as + `query`/`document` messages with the saved chat template. + + Args: + model: Local checkpoint path or Hugging Face model ID. + revision: Model revision to inspect. + + Returns: + The saved `processing_kwargs.chat_template` (possibly empty), or + `None` if the checkpoint does not declare the `message` modality. + + Raises: + ValueError: If the saved chat template settings are unsupported. + + """ + try: + st_config = get_hf_file_to_dict( + "config_sentence_transformers.json", model, revision + ) + if not isinstance(st_config, dict) or ( + st_config.get("model_type") != "CrossEncoder" + ): + return None + modules = get_hf_file_to_dict("modules.json", model, revision) + transformer_config = get_hf_file_to_dict( + "sentence_bert_config.json", model, revision + ) + except Exception: + logger.debug("Failed to read the Sentence Transformers config", exc_info=True) + return None + + if isinstance(modules, dict): + modules = modules.get("modules") + if not ( + isinstance(modules, list) + and len(modules) == 1 + and isinstance(modules[0], dict) + and modules[0].get("type") in _ST_TRANSFORMER_MODULE_TYPES + and not modules[0].get("path") + and isinstance(transformer_config, dict) + and transformer_config.get("transformer_task") == "sequence-classification" + and isinstance( + modality_config := transformer_config.get("modality_config"), dict + ) + and "message" in modality_config + ): + return None + + processing_kwargs = transformer_config.get("processing_kwargs") or {} + chat_template_kwargs = ( + processing_kwargs.get("chat_template") or {} + if isinstance(processing_kwargs, dict) + else None + ) + if not isinstance(chat_template_kwargs, dict): + raise ValueError( + "The Sentence Transformers processing_kwargs.chat_template must be " + "a JSON object." + ) + chat_template_kwargs = dict(chat_template_kwargs) + # vLLM does not restore the template suffix after truncation. + chat_template_kwargs.pop("restore_suffix", None) + if not isinstance(chat_template_kwargs.get("chat_template", ""), str): + raise ValueError( + "The Sentence Transformers chat template name must be a string." + ) + if reserved := {"tokenize", "tools"} & chat_template_kwargs.keys(): + raise ValueError( + f"Unsupported Sentence Transformers chat template settings: {reserved}" + ) + return chat_template_kwargs + + def _read_safetensors_metadata_in_dir(local_dir: Path) -> dict[str, Any]: return { param_name: info