diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index cdb5b00d4143..f122c423ba5c 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -723,7 +723,7 @@ steps: - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh" env: S3_BUCKET: "vllm-wheels" - VARIANT: "rocm721" + VARIANT: "rocm722" # ROCm Job 6: Build ROCm Release Docker Image - label: ":docker: Build release image - x86_64 - ROCm" diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 5940a4ee564d..a21916d0b531 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.1-complete +ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.2-complete ARG TRITON_BRANCH="ba5c1517" ARG TRITON_REPO="https://github.com/ROCm/triton.git" ARG PYTORCH_BRANCH="8514f051" # release/2.10 as of 3/17 @@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0" ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.10.post3" +ARG AITER_BRANCH="v0.1.12.post2" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" @@ -104,6 +104,28 @@ ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} +# torch profiler hotfix for 7.2.2: rebuild CLR with https://github.com/ROCm/rocm-systems/pull/5062 +# will be removed once we move to ROCm 7.2.3 +RUN apt-get update && apt-get install -y rocm-llvm-dev +RUN pip install CppHeaderParser +RUN git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems /tmp/rocm-systems \ + && cd /tmp/rocm-systems \ + && git sparse-checkout init --cone \ + && git sparse-checkout set projects/hip projects/clr \ + && git checkout 35e8c7bf8911862e5389509800e65fdf125412b3 \ + && export CLR_DIR=/tmp/rocm-systems/projects/clr \ + && export HIP_DIR=/tmp/rocm-systems/projects/hip \ + && mkdir -p $CLR_DIR/build && cd $CLR_DIR/build \ + && cmake \ + -DHIP_COMMON_DIR=$HIP_DIR \ + -DCMAKE_PREFIX_PATH="/opt/rocm/" \ + -DCLR_BUILD_HIP=ON \ + -DCLR_BUILD_OCL=OFF \ + -DHIP_PLATFORM=amd \ + .. \ + && make -j$(nproc) \ + && make install \ + && rm -rf /tmp/rocm-systems ### ### Triton Build @@ -153,8 +175,6 @@ RUN git clone ${PYTORCH_REPO} pytorch RUN cd pytorch && git checkout ${PYTORCH_BRANCH} RUN cd pytorch \ && pip install -r requirements.txt && git submodule update --init --recursive -RUN cd pytorch/third_party/kineto \ - && git remote add rocm https://github.com/ROCm/kineto && git fetch rocm && git checkout 2d73be3 RUN cd pytorch && python3 tools/amd_build/build_amd.py \ && if [ "$USE_SCCACHE" = "1" ]; then \ export HIP_CLANG_PATH=/opt/sccache-wrappers \ diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index dc4b5402cab6..83c5fc1b435f 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -155,6 +155,7 @@ Priority is **1 = highest** (tried first). | **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) | | **Head Sizes** | Supported attention head sizes | | **Sink** | Attention sink support (for StreamingLLM) | +| **Non-Causal** | Non-causal (bidirectional) attention support for decoder models | | **Sparse** | Sparse attention support (MLA only) | | **MM Prefix** | Multimodal prefix full attention support | | **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) | @@ -165,22 +166,22 @@ Priority is **1 = highest** (tried first). ## Standard Attention (MHA, MQA, GQA) Backends -| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ | -| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A | -| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x | -| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x | -| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 | -| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x | -| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 | -| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any | -| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any | -| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A | -| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A | -| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any | -| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any | -| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | Decoder | Any | +| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. | +| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ | +| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A | +| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x | +| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x | +| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 | +| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x | +| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 | +| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | +| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any | +| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | +| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A | +| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A | +| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | ❌ | Decoder | Any | +| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any | +| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`. > @@ -211,16 +212,16 @@ hardware and configuration. MLA decode backends are selected using the standard `-ac.backend=` argument (e.g., `FLASHMLA`, `TRITON_MLA`). -| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ | -| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x | -| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x | -| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | -| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | -| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | -| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | -| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | -| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any | +| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | +| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ | +| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | +| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x | +| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x | +| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | +| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | +| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | +| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | +| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | +| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | +| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any | +| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any | diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 804cd905e3b1..b23631484508 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -105,7 +105,7 @@ Batch invariance has been tested and verified on the following models: - **DeepSeek series**: `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1` - **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ` -- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct` +- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`, `Qwen/Qwen3-30B-A3B-Thinking-2507-FP8` - **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct` - **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct` - **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b` diff --git a/docs/serving/data_parallel_deployment.md b/docs/serving/data_parallel_deployment.md index 7b963b99d565..1f18b92f95b4 100644 --- a/docs/serving/data_parallel_deployment.md +++ b/docs/serving/data_parallel_deployment.md @@ -98,7 +98,7 @@ For larger scale deployments especially, it can make sense to handle the orchest In this case, it's more convenient to treat each DP rank like a separate vLLM deployment, with its own endpoint, and have an external router balance HTTP requests between them, making use of appropriate real-time telemetry from each server for routing decisions. -This can already be done trivially for non-MoE models, since each deployed server is fully independent. No data parallel CLI options need to be used for this. +This can already be done trivially for non-MoE models, since each deployed server is fully independent. In that case, launch independent vLLM instances without any `--data-parallel-*` arguments; external DP CLI options are only supported for MoE deployments. We support an equivalent topology for MoE DP+EP which can be configured via the following CLI arguments. diff --git a/mkdocs.yaml b/mkdocs.yaml index 6afc44d71af5..4b06b31ebe35 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -105,8 +105,7 @@ plugins: - https://docs.aiohttp.org/en/stable/objects.inv - https://pillow.readthedocs.io/en/stable/objects.inv - https://numpy.org/doc/stable/objects.inv - # TODO revert to stable once https://github.com/pytorch/pytorch/issues/182007 is fixed - - https://pytorch.org/docs/2.11/objects.inv + - https://pytorch.org/docs/stable/objects.inv - redirects: redirect_maps: features/spec_decode/README.md: features/speculative_decoding/README.md diff --git a/requirements/common.txt b/requirements/common.txt index 5d4519204ee9..652738eebe74 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -24,7 +24,7 @@ outlines_core == 0.2.14 # required for outlines backend disk cache diskcache == 5.6.3 lark == 1.2.2 -xgrammar >= 0.1.32, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le" +xgrammar >= 0.2.0, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le" typing_extensions >= 4.10 filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317 partial-json-parser # used for parsing partial JSON outputs diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 801af7db9db1..8445634ded40 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -42,6 +42,8 @@ anyio==4.13.0 # sse-starlette # starlette # watchfiles +apache-tvm-ffi==0.1.10 + # via xgrammar arctic-inference==0.1.1 # via -r requirements/test/rocm.in argcomplete==3.6.3 @@ -1264,6 +1266,7 @@ typing-extensions==4.15.0 # alembic # anthropic # anyio + # apache-tvm-ffi # azure-core # azure-identity # azure-storage-blob @@ -1345,7 +1348,7 @@ word2number==1.1 # via lm-eval wrapt==2.1.2 # via smart-open -xgrammar==0.1.33 +xgrammar==0.2.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/tests/benchmarks/test_custom_dataset_seed.py b/tests/benchmarks/test_custom_dataset_seed.py new file mode 100644 index 000000000000..dac87e6e6d98 --- /dev/null +++ b/tests/benchmarks/test_custom_dataset_seed.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import argparse +import json +from pathlib import Path + +import pytest +from transformers import AutoTokenizer, PreTrainedTokenizerBase + +from vllm.benchmarks.datasets import get_samples + + +@pytest.fixture(scope="session") +def hf_tokenizer() -> PreTrainedTokenizerBase: + return AutoTokenizer.from_pretrained("gpt2") + + +def _write_jsonl(path: Path, n_rows: int) -> None: + with path.open("w") as f: + for i in range(n_rows): + f.write(json.dumps({"prompt": f"row {i}: unique prompt content."}) + "\n") + + +def _args_for_custom(dataset_path: str, seed: int) -> argparse.Namespace: + return argparse.Namespace( + dataset_name="custom", + dataset_path=dataset_path, + disable_shuffle=False, + num_prompts=30, + custom_output_len=32, + skip_chat_template=True, + no_oversample=False, + seed=seed, + request_id_prefix="", + ) + + +@pytest.mark.benchmark +def test_custom_dataset_seed_propagates( + hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path +) -> None: + """--seed must control the CustomDataset shuffle used by get_samples. + + Without the fix, CustomDataset was instantiated without random_seed, + so its load-time shuffle always used DEFAULT_SEED=0 regardless of + args.seed, causing every run with --dataset-name custom to pick the + same subset of rows from a larger file. + """ + jsonl = tmp_path / "data.jsonl" + _write_jsonl(jsonl, n_rows=60) + + samples_a = get_samples(_args_for_custom(str(jsonl), seed=0), hf_tokenizer) + samples_b = get_samples(_args_for_custom(str(jsonl), seed=42), hf_tokenizer) + + prompts_a = {s.prompt for s in samples_a} + prompts_b = {s.prompt for s in samples_b} + + assert len(prompts_a) == 30 + assert len(prompts_b) == 30 + assert prompts_a != prompts_b + + +@pytest.mark.benchmark +def test_custom_dataset_same_seed_is_deterministic( + hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path +) -> None: + """Same --seed must yield the same CustomDataset subset.""" + jsonl = tmp_path / "data.jsonl" + _write_jsonl(jsonl, n_rows=60) + + samples_a = get_samples(_args_for_custom(str(jsonl), seed=7), hf_tokenizer) + samples_b = get_samples(_args_for_custom(str(jsonl), seed=7), hf_tokenizer) + + prompts_a = [s.prompt for s in samples_a] + prompts_b = [s.prompt for s in samples_b] + + assert prompts_a == prompts_b diff --git a/tests/conftest.py b/tests/conftest.py index 40adeda2bd50..779bd475f34b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -996,6 +996,8 @@ def generate( req_sample_output_ids: list[list[int]] = [] req_sample_output_strs: list[str] = [] req_logprobs = [] + if req_output.prompt_logprobs: + req_logprobs.extend(req_output.prompt_logprobs) for sample in req_output.outputs: output_str = sample.text output_ids = list(sample.token_ids) diff --git a/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py b/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py index a3df30fb02b2..fedbd74795b5 100644 --- a/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py +++ b/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py @@ -27,7 +27,8 @@ from ....utils import RemoteOpenAIServer # Tuned to prevent OOM on 18GB GPUs in transcription correctness tests. -MAX_SEQS_FOR_TRANSCRIPTION_TEST = 32 +MAX_SEQS_FOR_TRANSCRIPTION_TEST = 8 +GPU_UTIL_FOR_TRANSCRIPTION_TEST = 0.5 def to_bytes(y, sr): @@ -188,6 +189,7 @@ def test_wer_correctness( "--enforce-eager", f"--tokenizer_mode={model_info.tokenizer_mode}", f"--max_num_seqs={MAX_SEQS_FOR_TRANSCRIPTION_TEST}", + f"--gpu_memory_utilization={GPU_UTIL_FOR_TRANSCRIPTION_TEST}", ] if model_info.trust_remote_code: server_args.append("--trust-remote-code") diff --git a/tests/kernels/quantization/test_marlin_gemm.py b/tests/kernels/quantization/test_marlin_gemm.py index f918212f763c..390ad293ef70 100644 --- a/tests/kernels/quantization/test_marlin_gemm.py +++ b/tests/kernels/quantization/test_marlin_gemm.py @@ -6,17 +6,29 @@ """ import itertools +from types import SimpleNamespace import pytest import torch +import vllm.model_executor.kernels.linear.mixed_precision.marlin as marlin_module +import vllm.model_executor.parameter as parameter_module from tests.kernels.utils import opcheck from tests.quantization.utils import is_quant_method_supported from vllm import _custom_ops as ops +from vllm.model_executor.kernels.linear.mixed_precision.marlin import ( + MarlinLinearKernel, + _pad_parameter_output_dim, +) +from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( + MPLinearLayerConfig, +) from vllm.model_executor.layers.quantization.utils.int8_utils import ( per_token_quant_int8, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + GPTQ_MARLIN_MIN_THREAD_N, + check_marlin_supports_layer, marlin_make_empty_g_idx, marlin_make_workspace_new, marlin_permute_bias, @@ -42,6 +54,11 @@ quantize_weights, sort_weights, ) +from vllm.model_executor.parameter import ( + GroupQuantScaleParameter, + PackedColumnParameter, + PackedvLLMParameter, +) from vllm.platforms import current_platform from vllm.scalar_type import scalar_types @@ -618,3 +635,237 @@ def test_marlin_gemm_with_bias(size_m): max_diff = compute_max_diff(output, output_ref) assert max_diff < 0.04 + + +@pytest.mark.skipif( + not is_quant_method_supported("gptq_marlin"), + reason="Marlin is not supported on this GPU type.", +) +@pytest.mark.parametrize("orig_n", [32, 48, 96]) +def test_marlin_gemm_sub_tile_n_pad(orig_n): + """Check padded Marlin GEMM matches the unpadded reference output.""" + quant_type = scalar_types.uint4b8 + group_size = 128 + size_m, size_k = 32, 1024 + + padded_n = ( + (orig_n + GPTQ_MARLIN_MIN_THREAD_N - 1) // GPTQ_MARLIN_MIN_THREAD_N + ) * GPTQ_MARLIN_MIN_THREAD_N + + a_input = rand_data((size_m, size_k)) + b_weight = rand_data((size_k, orig_n)) + + b_weight_padded = torch.nn.functional.pad(b_weight, (0, padded_n - orig_n), value=0) + + w_ref_padded, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize( + b_weight_padded, quant_type, group_size, False + ) + + marlin_zp = marlin_make_empty_g_idx(marlin_s.device) + workspace = marlin_make_workspace_new(a_input.device) + + output_padded = ops.marlin_gemm( + a_input, + None, + marlin_q_w, + None, + marlin_s, + None, + None, + marlin_zp, + g_idx, + sort_indices, + workspace, + quant_type, + size_m, + padded_n, + size_k, + is_k_full=True, + use_atomic_add=False, + use_fp32_reduce=True, + is_zp_float=False, + ) + + output = output_padded[..., :orig_n].contiguous() + output_ref = torch.matmul(a_input, w_ref_padded[:, :orig_n]) + + torch.accelerator.synchronize() + + max_diff = compute_max_diff(output, output_ref) + assert max_diff < 0.04 + + +def test_marlin_supports_layer_uses_padded_output_n(): + layer = SimpleNamespace( + input_size=2048, + input_size_per_partition=2048, + output_size=64, + output_size_per_partition=32, + ) + + assert check_marlin_supports_layer(layer, group_size=128) + + unsupported_input_layer = SimpleNamespace( + input_size=2048, + input_size_per_partition=96, + output_size=64, + output_size_per_partition=32, + ) + + assert not check_marlin_supports_layer(unsupported_input_layer, group_size=128) + + +def _noop_weight_loader(*args, **kwargs): + pass + + +def test_marlin_output_padding_uses_parameter_layout(monkeypatch): + """Check output padding follows vLLM parameter layout metadata.""" + output_dim_pad = 32 + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + + gptq_qweight = PackedvLLMParameter( + data=torch.ones(2, 32, dtype=torch.int32), + input_dim=0, + output_dim=1, + packed_dim=0, + packed_factor=8, + weight_loader=_noop_weight_loader, + ) + _pad_parameter_output_dim(gptq_qweight, output_dim_pad) + assert gptq_qweight.shape == (2, 64) + assert torch.count_nonzero(gptq_qweight[:, 32:]) == 0 + + compressed_tensors_qweight = PackedvLLMParameter( + data=torch.ones(32, 2, dtype=torch.int32), + input_dim=1, + output_dim=0, + packed_dim=1, + packed_factor=8, + weight_loader=_noop_weight_loader, + ) + _pad_parameter_output_dim(compressed_tensors_qweight, output_dim_pad) + assert compressed_tensors_qweight.shape == (64, 2) + assert torch.count_nonzero(compressed_tensors_qweight[32:, :]) == 0 + + compressed_tensors_scales = GroupQuantScaleParameter( + data=torch.ones(32, 4, dtype=torch.float16), + input_dim=1, + output_dim=0, + weight_loader=_noop_weight_loader, + ) + _pad_parameter_output_dim(compressed_tensors_scales, output_dim_pad) + assert compressed_tensors_scales.shape == (64, 4) + assert torch.count_nonzero(compressed_tensors_scales[32:, :]) == 0 + + packed_qzeros = PackedColumnParameter( + data=torch.ones(4, 4, dtype=torch.int32), + output_dim=1, + packed_dim=1, + packed_factor=8, + weight_loader=_noop_weight_loader, + ) + _pad_parameter_output_dim(packed_qzeros, output_dim_pad) + assert packed_qzeros.shape == (4, 8) + assert torch.count_nonzero(packed_qzeros[:, 4:]) == 0 + + +@pytest.mark.skipif( + not is_quant_method_supported("gptq_marlin"), + reason="Marlin is not supported on this GPU type.", +) +def test_marlin_output_padding_keeps_deferred_bias_unpadded(monkeypatch): + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + + config = MPLinearLayerConfig( + full_weight_shape=(1024, 32), + partition_weight_shape=(1024, 32), + weight_type=scalar_types.uint4b8, + act_type=torch.float16, + group_size=128, + zero_points=False, + has_g_idx=False, + ) + kernel = MarlinLinearKernel(config, "qweight", "scales") + assert kernel.orig_output_size_per_partition == 32 + + layer = torch.nn.Module() + layer.qweight = PackedvLLMParameter( + data=torch.ones(128, 32, dtype=torch.int32), + input_dim=0, + output_dim=1, + packed_dim=0, + packed_factor=8, + weight_loader=_noop_weight_loader, + ) + layer.scales = GroupQuantScaleParameter( + data=torch.ones(8, 32, dtype=torch.float16), + input_dim=0, + output_dim=1, + weight_loader=_noop_weight_loader, + ) + layer.bias = torch.nn.Parameter(torch.ones(32, dtype=torch.float16)) + layer.g_idx_sort_indices = torch.empty(0, dtype=torch.int32) + + monkeypatch.setattr( + marlin_module, + "marlin_make_workspace_new", + lambda device: torch.empty(0, dtype=torch.int32), + ) + monkeypatch.setattr(kernel, "_transform_param", lambda layer, name, transform: None) + + permute_bias_shapes = [] + + def fake_marlin_permute_bias(bias): + permute_bias_shapes.append(tuple(bias.shape)) + return bias + + monkeypatch.setattr(marlin_module, "marlin_permute_bias", fake_marlin_permute_bias) + + kernel.process_weights_after_loading(layer) + assert layer.bias.shape == (32,) + assert permute_bias_shapes == [] + + captured_bias_shape = None + + def fake_apply_gptq_marlin_linear(**kwargs): + nonlocal captured_bias_shape + bias = kwargs["bias"] + captured_bias_shape = None if bias is None else tuple(bias.shape) + input_ = kwargs["input"] + return input_.new_zeros( + input_.shape[:-1] + (kwargs["output_size_per_partition"],) + ) + + monkeypatch.setattr( + marlin_module, "apply_gptq_marlin_linear", fake_apply_gptq_marlin_linear + ) + kernel.workspace = torch.empty(0, dtype=torch.int32) + kernel.is_k_full = True + + output = kernel.apply_weights( + layer, + torch.zeros(2, 1024, dtype=torch.float16), + layer.bias, + ) + + assert captured_bias_shape == (64,) + assert permute_bias_shapes == [(64,)] + assert output.shape == (2, 32) + assert layer.bias.shape == (32,) + + padded_bias = torch.ones(64, dtype=torch.float16) + output = kernel.apply_weights( + layer, + torch.zeros(2, 1024, dtype=torch.float16), + padded_bias, + ) + assert captured_bias_shape == (64,) + assert permute_bias_shapes == [(64,), (64,)] + assert output.shape == (2, 32) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 6e3e2d63e144..cf3553bd57de 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -59,6 +59,34 @@ def test_reload_lifecycle(): assert tensor.__dict__ == materialized_tensor.__dict__ +def test_materialize_layer_preserves_non_meta_tensors(): + """Ensure that materialize_layer does not overwrite non meta tensors.""" + layer = torch.nn.Linear(2, 3, bias=True) + + # Create a non meta bias tensor and meta weight, which can happen with FP8 + bias_values = torch.ones(3) + layer.bias.data.copy_(bias_values) + layer.weight = torch.nn.Parameter(layer.weight.data.to("meta")) + + assert layer.weight.is_meta + assert not layer.bias.is_meta + + # materialize the layer weights after the bias is initialized + info = LayerReloadingInfo( + restore_metadata=({}, {}), + restore_device=torch.device("cpu"), + ) + materialize_layer(layer, info) + + # Ensure the weight materialized off meta + assert not layer.weight.is_meta + assert layer.weight.device.type == "cpu" + + # Ensure that the bias is (still) not meta and values are unchanged + assert not layer.bias.is_meta + assert torch.equal(layer.bias.data, bias_values) + + def test_model_cleanup(dist_init, default_vllm_config): layer = QKVParallelLinear(2, 3, 4) assert layer.weight.weight_loader.__self__ is layer diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py new file mode 100644 index 000000000000..6922af79c08e --- /dev/null +++ b/tests/models/multimodal/test_nano_nemotron_vl.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.model_executor.models.nano_nemotron_vl import NemotronH_Nano_VL_V2 + + +class _TextOnlyMultiModalConfig: + def get_limit_per_prompt(self, modality: str) -> int: + return 0 + + +class _ImageOnlyMultiModalConfig: + def get_limit_per_prompt(self, modality: str) -> int: + return 1 if modality == "image" else 0 + + +class _ModelConfig: + multimodal_config = _TextOnlyMultiModalConfig() + + +class _ImageOnlyModelConfig: + multimodal_config = _ImageOnlyMultiModalConfig() + + +class _LanguageModel: + def __init__(self) -> None: + self.loaded_weights: list[tuple[str, object]] = [] + + def load_weights(self, weights): + self.loaded_weights = list(weights) + + +class _MissingMultiModalModule: + def named_parameters(self): + raise AssertionError("multimodal weights should not be inspected") + + def load_weights(self, weights): + raise AssertionError("multimodal weights should not be loaded") + + +class _AdapterModule: + def named_parameters(self): + return [] + + +class _VisionModel: + def __init__(self) -> None: + self.loaded_weights: list[tuple[str, object]] = [] + + def load_weights(self, weights): + self.loaded_weights = list(weights) + + +def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode(): + model = object.__new__(NemotronH_Nano_VL_V2) + language_model = _LanguageModel() + object.__setattr__(model, "model_config", _ModelConfig()) + object.__setattr__(model, "language_model", language_model) + object.__setattr__(model, "mlp1", _AdapterModule()) + object.__setattr__(model, "vision_model", _MissingMultiModalModule()) + object.__setattr__(model, "sound_encoder", None) + + language_weight = object() + model.load_weights( + [ + ("language_model.layers.0.weight", language_weight), + ("mlp1.0.weight", object()), + ("vision_model.radio_model.encoder.weight", object()), + ("sound_encoder.encoder.weight", object()), + ] + ) + + assert language_model.loaded_weights == [("layers.0.weight", language_weight)] + + +def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder(): + model = object.__new__(NemotronH_Nano_VL_V2) + language_model = _LanguageModel() + vision_model = _VisionModel() + object.__setattr__(model, "model_config", _ImageOnlyModelConfig()) + object.__setattr__(model, "language_model", language_model) + object.__setattr__(model, "mlp1", _AdapterModule()) + object.__setattr__(model, "vision_model", vision_model) + object.__setattr__(model, "sound_encoder", None) + + language_weight = object() + vision_weight = object() + model.load_weights( + [ + ("language_model.layers.0.weight", language_weight), + ("vision_model.radio_model.encoder.weight", vision_weight), + ] + ) + + assert language_model.loaded_weights == [("layers.0.weight", language_weight)] + assert vision_model.loaded_weights == [ + ("radio_model.encoder.weight", vision_weight) + ] + + +def test_nano_nemotron_vl_requires_sound_encoder_for_sound_weights(): + model = object.__new__(NemotronH_Nano_VL_V2) + language_model = _LanguageModel() + vision_model = _VisionModel() + object.__setattr__(model, "model_config", _ImageOnlyModelConfig()) + object.__setattr__(model, "language_model", language_model) + object.__setattr__(model, "mlp1", _AdapterModule()) + object.__setattr__(model, "vision_model", vision_model) + object.__setattr__(model, "sound_encoder", None) + + with pytest.raises(AssertionError): + model.load_weights([("sound_encoder.encoder.weight", object())]) diff --git a/tests/models/registry.py b/tests/models/registry.py index 21d3a50ce996..ec6c3473a785 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -946,13 +946,6 @@ def check_available_online( "HCXVisionForCausalLM": _HfExamplesInfo( "naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B", trust_remote_code=True, - max_transformers_version="4.57", - transformers_version_reason={ - "vllm": ( - "Custom config cannot be loaded with Transformers " - "v5 because `text_config` is not always set" - ) - }, ), "HCXVisionV2ForCausalLM": _HfExamplesInfo( "naver-hyperclovax/HyperCLOVAX-SEED-Think-32B", @@ -1148,30 +1141,17 @@ def check_available_online( "NemotronH_Nano_VL_V2": _HfExamplesInfo( "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", max_model_len=4096, - # NemotronH layers are constructed via `hybrid_override_pattern`: + # NemotronH layers are constructed via `hybrid_override_pattern` use_original_num_layers=True, hf_overrides={ - "vision_config": PretrainedConfig( - args={ - "min_num_patches": 1, # Trigger image dynamic res - "max_num_patches": 12, - "model": "vit_huge_patch16_224", - }, - # Trigger conv3d: - video_temporal_patch_size=2, - ), - "text_config": { - "num_hidden_layers": 2, - "hybrid_override_pattern": "M*", - }, + "text_config": {"num_hidden_layers": 2, "hybrid_override_pattern": "M*"}, }, trust_remote_code=True, ), - # NemotronH_Nano_Omni_Reasoning_V3 is an alias for NemotronH_Nano_VL_V2 - # Use the same registry test as NemotronH_Nano_VL_V2 above "NemotronH_Nano_Omni_Reasoning_V3": _HfExamplesInfo( - "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", + "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", max_model_len=4096, + # NemotronH layers are constructed via `hybrid_override_pattern` use_original_num_layers=True, hf_overrides={ "vision_config": PretrainedConfig( @@ -1181,35 +1161,17 @@ def check_available_online( "model": "vit_huge_patch16_224", }, video_temporal_patch_size=2, + # TODO(nhaber): This is `true` in the official `config.json`, + # but this causes a processor exception in the tests due to a known bug + # with mixed-resolution video when `true`. To be resolved. + video_maintain_aspect_ratio=False, ), - "text_config": { - "num_hidden_layers": 2, - "hybrid_override_pattern": "M*", - }, + "text_config": {"num_hidden_layers": 2, "hybrid_override_pattern": "M*"}, }, trust_remote_code=True, ), - # NemotronH_Super_Omni_Reasoning_V3 is an alias for NemotronH_Nano_VL_V2 as well - # Use the same registry test as NemotronH_Nano_VL_V2 above "NemotronH_Super_Omni_Reasoning_V3": _HfExamplesInfo( - "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", - max_model_len=4096, - use_original_num_layers=True, - hf_overrides={ - "vision_config": PretrainedConfig( - args={ - "min_num_patches": 1, - "max_num_patches": 12, - "model": "vit_huge_patch16_224", - }, - video_temporal_patch_size=2, - ), - "text_config": { - "num_hidden_layers": 2, - "hybrid_override_pattern": "M*", - }, - }, - trust_remote_code=True, + "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", is_available_online=False ), "OpenCUAForConditionalGeneration": _HfExamplesInfo( "xlangai/OpenCUA-7B", diff --git a/tests/quantization/test_cpu_offload.py b/tests/quantization/test_cpu_offload.py index 3b58614e58d4..151b5d97ddf3 100644 --- a/tests/quantization/test_cpu_offload.py +++ b/tests/quantization/test_cpu_offload.py @@ -70,4 +70,5 @@ def test_cpu_offload_compressed_tensors(monkeypatch): ["--enforce_eager"], ["--enforce_eager", "--cpu-offload-gb", "1"], max_wait_seconds=480, + include_seeded_sampling=False, ) diff --git a/tests/quantization/test_turboquant.py b/tests/quantization/test_turboquant.py index f074ce119ae8..b9567195b3a8 100644 --- a/tests/quantization/test_turboquant.py +++ b/tests/quantization/test_turboquant.py @@ -182,22 +182,100 @@ def test_all_presets_all_head_dims(self, preset, head_dim): # ---- Boundary skip layers ---- + @staticmethod + def _dense_model_config(num_layers): + from types import SimpleNamespace + + return SimpleNamespace( + is_hybrid=False, + hf_text_config=SimpleNamespace(num_hidden_layers=num_layers), + ) + def test_boundary_skip_layers_basic(self): - layers = TurboQuantConfig.get_boundary_skip_layers(32) + mc = self._dense_model_config(32) + layers = TurboQuantConfig.get_boundary_skip_layers(mc) assert layers == ["0", "1", "30", "31"] def test_boundary_skip_layers_zero(self): - assert TurboQuantConfig.get_boundary_skip_layers(32, 0) == [] + mc = self._dense_model_config(32) + assert TurboQuantConfig.get_boundary_skip_layers(mc, 0) == [] def test_boundary_skip_layers_small_model(self): - layers = TurboQuantConfig.get_boundary_skip_layers(4) + mc = self._dense_model_config(4) + layers = TurboQuantConfig.get_boundary_skip_layers(mc) assert layers == ["0", "1", "2", "3"] def test_boundary_skip_layers_cap_at_half(self): - layers = TurboQuantConfig.get_boundary_skip_layers(8, 10) + mc = self._dense_model_config(8) + layers = TurboQuantConfig.get_boundary_skip_layers(mc, 10) assert len(layers) == 8 +class TestHybridAttentionIndices: + """Regression tests for boundary protection on hybrid models. + + Hybrid models (attention + Mamba / linear-attention) identify KV-carrying + layers via layer_types / layers_block_type / attn_type_list. The helper + must return the *global* layer indices of the full-attention layers so + that kv_cache_dtype_skip_layers matches what extract_layer_index(prefix) + reports on the Attention layers at runtime. + """ + + @staticmethod + def _fake_model_config(text_cfg=None, hf_cfg=None): + from types import SimpleNamespace + + return SimpleNamespace( + hf_text_config=text_cfg if text_cfg is not None else SimpleNamespace(), + hf_config=hf_cfg if hf_cfg is not None else SimpleNamespace(), + ) + + def test_layer_types_full_attention(self): + from vllm.model_executor.layers.quantization.turboquant.config import ( + _get_full_attention_layer_indices, + ) + + cfg = type("C", (), {})() + cfg.layer_types = [ + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "full_attention", + "full_attention", + ] + mc = self._fake_model_config(text_cfg=cfg) + assert _get_full_attention_layer_indices(mc) == [2, 4, 5] + + def test_layers_block_type_jamba(self): + from vllm.model_executor.layers.quantization.turboquant.config import ( + _get_full_attention_layer_indices, + ) + + cfg = type("C", (), {})() + cfg.layers_block_type = ["mamba", "attention", "mamba", "attention"] + mc = self._fake_model_config(text_cfg=cfg) + assert _get_full_attention_layer_indices(mc) == [1, 3] + + def test_attn_type_list_minimax(self): + from vllm.model_executor.layers.quantization.turboquant.config import ( + _get_full_attention_layer_indices, + ) + + hf = type("C", (), {})() + hf.attn_type_list = [0, 1, 0, 1, 1] + mc = self._fake_model_config(hf_cfg=hf) + assert _get_full_attention_layer_indices(mc) == [1, 3, 4] + + def test_no_hybrid_hints_returns_empty(self): + from vllm.model_executor.layers.quantization.turboquant.config import ( + _get_full_attention_layer_indices, + ) + + mc = self._fake_model_config() + assert _get_full_attention_layer_indices(mc) == [] + + # ============================================================================ # Centroids tests (CPU-only) # ============================================================================ diff --git a/tests/reasoning/test_kimi_k2_reasoning_parser.py b/tests/reasoning/test_kimi_k2_reasoning_parser.py index 0f80bb8854a8..dfce2075c6a9 100644 --- a/tests/reasoning/test_kimi_k2_reasoning_parser.py +++ b/tests/reasoning/test_kimi_k2_reasoning_parser.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock + import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest @@ -12,6 +14,20 @@ REASONING_MODEL_NAME = "moonshotai/Kimi-K2.5" +@pytest.fixture +def mock_kimi_k2_tokenizer(): + tokenizer = MagicMock() + tokenizer.get_vocab.return_value = { + "": 100, + "": 101, + "<|tool_calls_section_begin|>": 200, + "<|tool_calls_section_end|>": 201, + "<|tool_call_begin|>": 202, + "<|tool_call_end|>": 203, + } + return tokenizer + + @pytest.fixture(scope="module") def kimi_k2_tokenizer(): return get_tokenizer(tokenizer_name=REASONING_MODEL_NAME, trust_remote_code=True) @@ -153,3 +169,50 @@ def test_streaming_tool_section_ends_reasoning(kimi_k2_tokenizer): ) assert isinstance(result, DeltaMessage) assert result.content == "<|tool_calls_section_begin|>" + + +def test_streaming_end_token_id_buffered(mock_kimi_k2_tokenizer): + """When stop sequences buffer text, ID arrives before its text. + + The token ID is present in delta_token_ids but the actual string is not + yet in delta_text (still buffered). The parser must return None to wait + for the next delta, instead of calling find() which returns -1 and + silently corrupting the text split. + """ + parser = KimiK2ReasoningParser(mock_kimi_k2_tokenizer) + think_id = parser._start_token_id + end_think_id = parser._end_token_id + + # Simulate: ID arrived but text not yet flushed. + # Two token IDs in delta to bypass the single-special-token guard. + result = parser.extract_reasoning_streaming( + previous_text="some reasoning", + current_text="some reasoning extra", + delta_text="extra", # text not yet flushed + previous_token_ids=[think_id], + current_token_ids=[think_id, end_think_id, 999], + delta_token_ids=[end_think_id, 999], + ) + assert result is None + + +def test_streaming_tool_section_id_buffered(mock_kimi_k2_tokenizer): + """When stop sequences buffer text, tool section start ID arrives before its text. + + Same buffering scenario as above but for <|tool_calls_section_begin|>. + Without the guard, find() returns -1 and delta_text[:tool_index] silently + drops the last character of reasoning. + """ + parser = KimiK2ReasoningParser(mock_kimi_k2_tokenizer) + think_id = parser._start_token_id + tool_begin_id = parser._tool_section_start_token_id + + result = parser.extract_reasoning_streaming( + previous_text="some reasoning", + current_text="some reasoning extra", + delta_text="extra", # tool section text not yet flushed + previous_token_ids=[think_id], + current_token_ids=[think_id, tool_begin_id, 999], + delta_token_ids=[tool_begin_id, 999], + ) + assert result is None diff --git a/tests/test_config.py b/tests/test_config.py index 02e4d1d5d77b..57d1e1bc686b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1215,8 +1215,6 @@ def test_scheduler_config_init(): ("facebook/opt-125m", 1, False, False), # Non-MoE model with DP>1 internal LB should need coordinator ("facebook/opt-125m", 2, False, True), - # Non-MoE model with DP>1 external LB should not need coordinator - ("facebook/opt-125m", 2, True, False), # MoE model with DP=1 should not need coordinator ("mistralai/Mixtral-8x7B-Instruct-v0.1", 1, False, False), # MoE model with DP>1 internal LB should need both coordinator diff --git a/tests/tool_parsers/test_deepseekv4_tool_parser.py b/tests/tool_parsers/test_deepseekv4_tool_parser.py index 631d0fb97b33..cc77a1f77756 100644 --- a/tests/tool_parsers/test_deepseekv4_tool_parser.py +++ b/tests/tool_parsers/test_deepseekv4_tool_parser.py @@ -6,6 +6,15 @@ import json from unittest.mock import MagicMock +import pytest +from xgrammar import StructuralTag + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedFunction, + ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, + ChatCompletionToolsParam, +) from vllm.tool_parsers import ToolParserManager from vllm.tool_parsers.deepseekv4_tool_parser import DeepSeekV4ToolParser @@ -20,6 +29,43 @@ PARAM_END = "" +@pytest.fixture +def sample_tools() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_current_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "The city name"}, + "state": {"type": "string", "description": "The state code"}, + "unit": {"type": "string", "enum": ["fahrenheit", "celsius"]}, + }, + "required": ["city", "state"], + }, + }, + ), + ChatCompletionToolsParam( + type="function", + function={ + "name": "calculate_area", + "description": "Calculate area of a shape", + "parameters": { + "type": "object", + "properties": { + "shape": {"type": "string"}, + "dimensions": {"type": "object"}, + "precision": {"type": "integer"}, + }, + }, + }, + ), + ] + + def make_parser(tools=None) -> DeepSeekV4ToolParser: return DeepSeekV4ToolParser(MOCK_TOKENIZER, tools=tools) @@ -121,3 +167,39 @@ def test_streaming_extracts_complete_invokes(): ] assert names == ["search"] assert json.loads(reconstruct_args(deltas)) == {"query": "deepseek v4"} + + +def test_get_vllm_registry_structural_tag_returns_structural_tag( + sample_tools: list[ChatCompletionToolsParam], +) -> None: + parser = make_parser() + req = ChatCompletionRequest( + messages=[], + model="m", + tools=sample_tools, + tool_choice="auto", + ) + tag = parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) + + req = ChatCompletionRequest( + messages=[], + model="m", + tools=sample_tools, + tool_choice="required", + ) + tag = parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) + + if sample_tools: + tool = sample_tools[0] + req = ChatCompletionRequest( + messages=[], + model="m", + tools=sample_tools, + ) + req.tool_choice = ChatCompletionNamedToolChoiceParam( + function=ChatCompletionNamedFunction(name=tool.function.name) + ) + tag = parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) diff --git a/tests/tool_parsers/test_qwen3coder_tool_parser.py b/tests/tool_parsers/test_qwen3coder_tool_parser.py index c62e95830243..26bbf1a044bc 100644 --- a/tests/tool_parsers/test_qwen3coder_tool_parser.py +++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py @@ -6,8 +6,11 @@ import pytest from openai.types.responses.function_tool import FunctionTool +from xgrammar import StructuralTag from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedFunction, + ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionToolsParam, ) @@ -108,6 +111,27 @@ def sample_tools(request): ] +def _as_chat_completion_tools( + tools: list[ChatCompletionToolsParam | FunctionTool], +) -> list[ChatCompletionToolsParam]: + normalized: list[ChatCompletionToolsParam] = [] + for tool in tools: + if isinstance(tool, ChatCompletionToolsParam): + normalized.append(tool) + else: + normalized.append( + ChatCompletionToolsParam( + type="function", + function={ + "name": tool.name, + "description": tool.description, + "parameters": tool.parameters, + }, + ) + ) + return normalized + + def assert_tool_calls( actual_tool_calls: list[ToolCall], expected_tool_calls: list[ToolCall] ): @@ -1146,3 +1170,88 @@ def test_no_double_serialization_string_args(qwen3_tool_parser): args = json.loads(raw_arguments) assert args["message"] == "hello world" assert '\\"hello world\\"' not in raw_arguments + + +def test_get_vllm_registry_structural_tag_returns_structural_tag( + qwen3_tool_parser: Qwen3CoderToolParser, + sample_tools: list[ChatCompletionToolsParam], +) -> None: + request_tools = _as_chat_completion_tools(sample_tools) + req = ChatCompletionRequest( + messages=[], + model="m", + tools=request_tools, + tool_choice="auto", + ) + tag = qwen3_tool_parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) + + req = ChatCompletionRequest( + messages=[], + model="m", + tools=request_tools, + tool_choice="required", + ) + tag = qwen3_tool_parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) + + if request_tools: + tool = request_tools[0] + req = ChatCompletionRequest( + messages=[], + model="m", + tools=request_tools, + ) + req.tool_choice = ChatCompletionNamedToolChoiceParam( + function=ChatCompletionNamedFunction(name=tool.function.name) + ) + tag = qwen3_tool_parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) + + +@pytest.mark.parametrize("include_reasoning", [True, False]) +def test_adjust_request_auto_uses_vllm_registry_structural_tag( + monkeypatch: pytest.MonkeyPatch, + qwen3_tool_parser: Qwen3CoderToolParser, + sample_tools: list[ChatCompletionToolsParam], + include_reasoning: bool, +) -> None: + monkeypatch.setattr( + "vllm.tool_parsers.abstract_tool_parser.VLLM_ENFORCE_STRICT_TOOL_CALLING", + True, + ) + request_tools = _as_chat_completion_tools(sample_tools) + req = ChatCompletionRequest( + messages=[], + model="m", + tools=request_tools, + tool_choice="auto", + include_reasoning=include_reasoning, + ) + out = qwen3_tool_parser.adjust_request(req) + assert out.structured_outputs is not None + assert out.structured_outputs.structural_tag is not None + assert isinstance(out.structured_outputs.structural_tag, str) + loaded = json.loads(out.structured_outputs.structural_tag) + assert isinstance(loaded, dict) + + +def test_adjust_request_required_prefers_structural_tag( + monkeypatch: pytest.MonkeyPatch, + qwen3_tool_parser: Qwen3CoderToolParser, + sample_tools: list[ChatCompletionToolsParam], +) -> None: + monkeypatch.setattr( + "vllm.tool_parsers.abstract_tool_parser.VLLM_ENFORCE_STRICT_TOOL_CALLING", + True, + ) + request_tools = _as_chat_completion_tools(sample_tools) + req = ChatCompletionRequest( + messages=[], + model="m", + tools=request_tools, + tool_choice="required", + ) + out = qwen3_tool_parser.adjust_request(req) + assert out.structured_outputs is not None + assert out.structured_outputs.structural_tag is not None diff --git a/tests/utils.py b/tests/utils.py index e4b6a6ff6e70..41202aa19481 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +import atexit import contextlib import copy import functools @@ -134,6 +135,11 @@ class RemoteVLLMServer: """ DUMMY_API_KEY = "token-abc123" # vLLM's OpenAI server does not need API key + _active_servers: set["RemoteVLLMServer"] = set() + _active_servers_lock = threading.RLock() + _cleanup_hooks_registered = False + _signal_hooks_registered = False + _previous_signal_handlers: dict[int, Any] = {} proc: subprocess.Popen def _create_cli_subcommand(self): @@ -209,6 +215,7 @@ def __init__( ) self._pre_download_model(model, args) + self._shutdown_complete = False # Record GPU memory before server start so we know what # "released" looks like. @@ -221,6 +228,7 @@ def __init__( ) self._start_server(model, vllm_serve_args, env_dict) + self._register_active_server() max_wait_seconds = max_wait_seconds or 480 try: self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds) @@ -246,8 +254,70 @@ def _shutdown(self) -> None: (when the server fails to start). Must be safe to call even if the process is already dead. """ - self._terminate_process_tree() - self._wait_for_gpu_memory_release() + if self._shutdown_complete: + return + + self._shutdown_complete = True + try: + self._terminate_process_tree() + self._wait_for_gpu_memory_release() + finally: + self._unregister_active_server() + + @classmethod + def _ensure_cleanup_hooks_registered(cls) -> None: + """Register process-exit cleanup for detached server subprocesses.""" + root_cls = RemoteVLLMServer + with root_cls._active_servers_lock: + if not root_cls._cleanup_hooks_registered: + atexit.register(root_cls._shutdown_active_servers) + root_cls._cleanup_hooks_registered = True + + if ( + threading.current_thread() is threading.main_thread() + and not root_cls._signal_hooks_registered + ): + for signum in (signal.SIGTERM, signal.SIGINT): + root_cls._previous_signal_handlers[signum] = signal.getsignal( + signum + ) + signal.signal(signum, root_cls._handle_parent_signal) + root_cls._signal_hooks_registered = True + + def _register_active_server(self) -> None: + """Track this server so parent-process exits still clean it up.""" + RemoteVLLMServer._ensure_cleanup_hooks_registered() + with RemoteVLLMServer._active_servers_lock: + RemoteVLLMServer._active_servers.add(self) + + def _unregister_active_server(self) -> None: + with RemoteVLLMServer._active_servers_lock: + RemoteVLLMServer._active_servers.discard(self) + + @classmethod + def _shutdown_active_servers(cls) -> None: + """Best-effort shutdown for all live RemoteVLLMServer instances.""" + with cls._active_servers_lock: + servers = list(cls._active_servers) + + for server in servers: + with contextlib.suppress(Exception): + server._shutdown() + + @classmethod + def _handle_parent_signal(cls, signum, frame) -> None: + """Clean up detached servers before letting the signal terminate pytest.""" + cls._shutdown_active_servers() + + previous_handler = cls._previous_signal_handlers.get(signum, signal.SIG_DFL) + if callable(previous_handler): + previous_handler(signum, frame) + elif previous_handler == signal.SIG_IGN: + return + elif signum == signal.SIGINT: + raise KeyboardInterrupt + else: + raise SystemExit(128 + signum) def _terminate_process_tree(self) -> None: """Kill the server process tree without waiting for GPU memory release. @@ -315,6 +385,9 @@ def shutdown_many(cls, servers: Sequence["RemoteVLLMServer"]) -> None: if not servers: return + for server in servers: + server._shutdown_complete = True + threads = [ threading.Thread( target=s._terminate_process_tree, @@ -339,7 +412,11 @@ def shutdown_many(cls, servers: Sequence["RemoteVLLMServer"]) -> None: else s._pre_server_gpu_memory ), ) - earliest._wait_for_gpu_memory_release() + try: + earliest._wait_for_gpu_memory_release() + finally: + for server in servers: + server._unregister_active_server() def _kill_process_group_survivors( self, pgid: int | None, timeout: float = 15.0 @@ -705,6 +782,7 @@ def _test_completion( model: str, prompt: str, token_ids: list[int], + include_seeded_sampling: bool = True, ): results = [] @@ -739,33 +817,40 @@ def _test_completion( } ) - # test seeded random sampling - completion = client.completions.create( - model=model, prompt=prompt, max_tokens=5, seed=33, temperature=1.0 - ) + if include_seeded_sampling: + # test seeded random sampling + completion = client.completions.create( + model=model, prompt=prompt, max_tokens=5, seed=33, temperature=1.0 + ) - results.append( - { - "test": "seeded_sampling", - "text": completion.choices[0].text, - "finish_reason": completion.choices[0].finish_reason, - "usage": completion.usage, - } - ) + results.append( + { + "test": "seeded_sampling", + "text": completion.choices[0].text, + "finish_reason": completion.choices[0].finish_reason, + "usage": completion.usage, + } + ) - # test seeded random sampling with multiple prompts - completion = client.completions.create( - model=model, prompt=[prompt, prompt], max_tokens=5, seed=33, temperature=1.0 - ) + # test seeded random sampling with multiple prompts + completion = client.completions.create( + model=model, + prompt=[prompt, prompt], + max_tokens=5, + seed=33, + temperature=1.0, + ) - results.append( - { - "test": "seeded_sampling", - "text": [choice.text for choice in completion.choices], - "finish_reason": [choice.finish_reason for choice in completion.choices], - "usage": completion.usage, - } - ) + results.append( + { + "test": "seeded_sampling", + "text": [choice.text for choice in completion.choices], + "finish_reason": [ + choice.finish_reason for choice in completion.choices + ], + "usage": completion.usage, + } + ) # test simple list batch = client.completions.create( @@ -960,6 +1045,7 @@ def compare_two_settings( *, method: str = "generate", max_wait_seconds: float | None = None, + include_seeded_sampling: bool = True, ) -> None: """ Launch API server with two different sets of arguments/environments @@ -971,6 +1057,8 @@ def compare_two_settings( arg2: The second set of arguments to pass to the API server. env1: The first set of environment variables to pass to the API server. env2: The second set of environment variables to pass to the API server. + include_seeded_sampling: Whether to include temperature=1.0 seeded + sampling checks in the default generate comparison. """ compare_all_settings( @@ -979,6 +1067,7 @@ def compare_two_settings( [env1, env2], method=method, max_wait_seconds=max_wait_seconds, + include_seeded_sampling=include_seeded_sampling, ) @@ -989,6 +1078,7 @@ def compare_all_settings( *, method: str = "generate", max_wait_seconds: float | None = None, + include_seeded_sampling: bool = True, ) -> None: """ Launch API server with several different sets of arguments/environments @@ -997,6 +1087,8 @@ def compare_all_settings( model: The model to test. all_args: A list of argument lists to pass to the API server. all_envs: A list of environment dictionaries to pass to the API server. + include_seeded_sampling: Whether to include temperature=1.0 seeded + sampling checks in the default generate comparison. """ trust_remote_code = False @@ -1057,7 +1149,13 @@ def compare_all_settings( ) if method == "generate": - results += _test_completion(client, model, prompt, token_ids) + results += _test_completion( + client, + model, + prompt, + token_ids, + include_seeded_sampling=include_seeded_sampling, + ) elif method == "generate_close": results += _test_completion_close(client, model, prompt) elif method == "generate_chat": diff --git a/tests/v1/attention/test_kv_head_stride_canonicalization.py b/tests/v1/attention/test_kv_head_stride_canonicalization.py new file mode 100644 index 000000000000..635f46390cfc --- /dev/null +++ b/tests/v1/attention/test_kv_head_stride_canonicalization.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for canonicalize_singleton_dim_strides. + +Background +---------- +When num_kv_heads_per_rank == 1 (e.g. Qwen3.5-397B with TP=8 → 1 KV head +per rank), PyTorch's is_contiguous() returns True for *any* stride on the +size-1 dimension. The KV cache allocator can therefore produce a tensor +where that singleton dim has stride = 1 element (2 bytes for bf16) instead +of the canonical product-of-remaining-dims value. + +CUDA TMA (used by FlashInfer XQA SM90 and Flash-Attention 3/4 on H100+) +requires all non-outermost strides to be multiples of 16 bytes. A 2-byte +stride triggers cudaErrorIllegalInstruction. + +canonicalize_singleton_dim_strides() patches degenerate strides on all +size-1 dimensions via torch.as_strided — zero-copy. + +The degenerate stride manifests at different positions in different backends: +- FlashInfer: stride(-3) after kv_cache.permute() → shape [..., 1, B, D] +- FlashAttention: stride(-2) after kv_cache.unbind(0) → shape [N, B, 1, D] +""" + +import torch + +from vllm.utils.torch_utils import canonicalize_singleton_dim_strides + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _inject_degenerate_stride(t: torch.Tensor, dim: int) -> torch.Tensor: + """Return a view of t with a degenerate (stride=1) on a size-1 dim.""" + assert t.shape[dim] == 1, f"dim {dim} must have size 1" + strides = list(t.stride()) + strides[dim] = 1 # inject the bug + return t.as_strided(t.shape, strides) + + +# --------------------------------------------------------------------------- +# Tests: canonicalize_singleton_dim_strides +# --------------------------------------------------------------------------- + + +class TestCanonicalizeSingletonDimStrides: + def test_flashinfer_layout_dim_neg3(self): + """FlashInfer path: degenerate stride at dim -3 (num_kv_heads).""" + # Shape after permute: [num_blocks, 2, num_kv_heads, block_size, head_size] + num_blocks, block_size, head_size = 64, 16, 128 + t = torch.zeros(num_blocks, 2, 1, block_size, head_size, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-3) + + assert t_deg.stride(-3) == 1 # confirm degenerate + assert t_deg.is_contiguous() # PyTorch doesn't notice + + fixed = canonicalize_singleton_dim_strides(t_deg) + + assert fixed.stride(-3) == block_size * head_size # canonical = 2048 + assert fixed.stride(-2) == head_size # inner dims unchanged + assert fixed.stride(-1) == 1 + + def test_flash_attn_layout_dim_neg2(self): + """FlashAttention path: degenerate stride at dim -2 (num_kv_heads).""" + # Shape after unbind(0): [num_blocks, block_size, num_kv_heads, head_size] + num_blocks, block_size, head_size = 64, 16, 128 + t = torch.zeros(num_blocks, block_size, 1, head_size, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-2) + + assert t_deg.stride(-2) == 1 + assert t_deg.is_contiguous() + + fixed = canonicalize_singleton_dim_strides(t_deg) + + assert fixed.stride(-2) == head_size # canonical = 128 + assert fixed.stride(-1) == 1 + + def test_canonical_strides_returned_as_is(self): + """No degenerate strides → same object returned (no copy, no new view).""" + t = torch.zeros(64, 2, 1, 16, 128, dtype=torch.bfloat16) + result = canonicalize_singleton_dim_strides(t) + assert result is t + + def test_multi_kv_heads_unchanged(self): + """num_kv_heads > 1 → strides are already canonical → unchanged.""" + t = torch.zeros(16, 2, 4, 16, 128, dtype=torch.bfloat16) + original_strides = t.stride() + result = canonicalize_singleton_dim_strides(t) + assert result.stride() == original_strides + + def test_data_pointer_preserved(self): + """Fix is zero-copy: same underlying storage.""" + t = torch.zeros(8, 2, 1, 16, 128, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-3) + fixed = canonicalize_singleton_dim_strides(t_deg) + assert fixed.data_ptr() == t_deg.data_ptr() + assert fixed.storage_offset() == t_deg.storage_offset() + + def test_multiple_singleton_dims(self): + """All size-1 dims with degenerate strides are fixed.""" + # Shape: [1, 1, 8, 32] — two size-1 dims + t = torch.zeros(1, 1, 8, 32, dtype=torch.float16) + # Both size-1 dims get degenerate strides + t_deg = t.as_strided(t.shape, (1, 1, 32, 1)) # both leading dims = 1 + + fixed = canonicalize_singleton_dim_strides(t_deg) + + assert fixed.stride(0) == 1 * 8 * 32 # canonical: 256 + assert fixed.stride(1) == 1 * 8 * 32 # canonical: 256 (same since size-1) + assert fixed.stride(2) == 32 + assert fixed.stride(3) == 1 + + def test_various_shapes_flashinfer(self): + """Correctness across different block_size / head_size for FlashInfer layout.""" + for block_size, head_size in [(16, 64), (16, 128), (32, 128), (16, 256)]: + t = torch.zeros(8, 2, 1, block_size, head_size, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-3) + fixed = canonicalize_singleton_dim_strides(t_deg) + assert fixed.stride(-3) == block_size * head_size, ( + f"Failed for block_size={block_size}, head_size={head_size}: " + f"got stride(-3)={fixed.stride(-3)}" + ) + + def test_various_shapes_flash_attn(self): + """Correctness across different shapes for FlashAttention layout.""" + for block_size, head_size in [(16, 64), (16, 128), (32, 128)]: + t = torch.zeros(8, block_size, 1, head_size, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-2) + fixed = canonicalize_singleton_dim_strides(t_deg) + assert fixed.stride(-2) == head_size, ( + f"Failed for block_size={block_size}, head_size={head_size}: " + f"got stride(-2)={fixed.stride(-2)}" + ) + + def test_tma_alignment_satisfied_after_fix_bf16(self): + """After fix, all strides meet 16-byte TMA alignment for bf16.""" + t = torch.zeros(64, 2, 1, 16, 128, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-3) + fixed = canonicalize_singleton_dim_strides(t_deg) + + element_size = fixed.element_size() # 2 bytes for bf16 + for i, s in enumerate(fixed.stride()): + assert (s * element_size) % 16 == 0 or i == len(fixed.stride()) - 1, ( + f"dim {i} stride {s} * {element_size} bytes not 16-byte aligned" + ) + + def test_non_contiguous_outer_dims_preserved(self): + """Outer (non-size-1) non-contiguous strides are left unchanged.""" + # Simulate cross-layer unified allocation: num_blocks stride is non-canonical + # but the inner dims should be fixed. + base = torch.zeros(200, 2, 1, 16, 128, dtype=torch.bfloat16) + # Slice every 2nd block → non-canonical outer stride + t_sliced = base[::2] # shape [100, 2, 1, 16, 128], stride[0] = 2*canonical + t_deg = _inject_degenerate_stride(t_sliced, dim=-3) + + fixed = canonicalize_singleton_dim_strides(t_deg) + + # Outer stride should be unchanged (not a size-1 dim) + assert fixed.stride(0) == t_sliced.stride(0) + # Inner degenerate stride should be fixed + assert fixed.stride(-3) == 16 * 128 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index f91ea85779d5..8ab47b618957 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers.attention.mla_attention import ( QueryLenSupport, _DecodeConcatQuantFP8, + get_mla_prefill_scale, ) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape @@ -785,7 +786,8 @@ def test_backend_correctness( assert kv_lora_rank + qk_rope_head_dim == head_size, ( f"MLA dimensions don't match: {total_head_size} != {head_size}" ) - scale = 1.0 / (total_head_size**0.5) + decode_scale = 1.0 / (total_head_size**0.5) + prefill_scale = get_mla_prefill_scale(vllm_config.model_config) # 2. Generate data and compute SDPA reference output for MLA all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], [] @@ -902,7 +904,7 @@ def test_backend_correctness( v_sdpa_in = v_mqa.unsqueeze(0).transpose(1, 2) sdpa_out_i_decode = torch.nn.functional.scaled_dot_product_attention( - q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale + q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=decode_scale ) sdpa_out_i_decode = sdpa_out_i_decode.transpose(1, 2).squeeze( 0 @@ -938,7 +940,7 @@ def test_backend_correctness( # Single attention call with custom mask sdpa_out_i_prefill = torch.nn.functional.scaled_dot_product_attention( - q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale + q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=prefill_scale ) sdpa_out_i_prefill = sdpa_out_i_prefill.transpose(1, 2).squeeze(0) sdpa_out_i_prefill = sdpa_out_i_prefill.flatten(start_dim=-2) diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index 068eb43faf40..873cfb18701b 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -2,12 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for MLA prefill backend selector.""" +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest import torch from vllm.config import AttentionConfig, ModelConfig, VllmConfig +from vllm.model_executor.layers.attention.mla_attention import get_mla_prefill_scale +from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import ( + yarn_get_mscale, +) from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.mla.prefill.selector import ( @@ -53,6 +58,62 @@ def _make_vllm_config( return mock_vllm_config +class TestMLAPrefillScale: + """Tests for the MLA prefill softmax scale.""" + + def test_uses_qk_head_dim_for_deepseek_v2_style_mla(self): + model_config = SimpleNamespace( + hf_text_config=SimpleNamespace( + q_lora_rank=None, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + rope_parameters={"rope_type": "default"}, + ) + ) + + assert get_mla_prefill_scale(model_config) == pytest.approx(192**-0.5) + + def test_applies_deepseek_yarn_mscale(self): + model_config = SimpleNamespace( + hf_text_config=SimpleNamespace( + q_lora_rank=None, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + rope_parameters={ + "rope_type": "yarn", + "factor": 40, + "mscale_all_dim": 0.707, + }, + ) + ) + + mscale = yarn_get_mscale(40, 0.707) + assert get_mla_prefill_scale(model_config) == pytest.approx( + 192**-0.5 * mscale * mscale + ) + + def test_deepseek_v4_style_mla_does_not_apply_yarn_mscale(self): + model_config = SimpleNamespace( + hf_text_config=SimpleNamespace( + compress_ratios=[4], + q_lora_rank=1536, + head_dim=128, + qk_rope_head_dim=64, + rope_parameters={ + "rope_type": "yarn", + "factor": 40, + "mscale_all_dim": 0.707, + }, + ) + ) + + assert get_mla_prefill_scale(model_config) == pytest.approx(128**-0.5) + + class TestGetMLAPrefillBackend: """Tests for get_mla_prefill_backend (public API).""" diff --git a/tests/v1/distributed/test_external_lb_dp.py b/tests/v1/distributed/test_external_lb_dp.py index cfef8449ebf8..06e8e574a05d 100644 --- a/tests/v1/distributed/test_external_lb_dp.py +++ b/tests/v1/distributed/test_external_lb_dp.py @@ -14,7 +14,7 @@ from tests.utils import RemoteOpenAIServer from vllm.platforms import current_platform -MODEL_NAME = "ibm-research/PowerMoE-3b" +MODEL_NAME = os.getenv("MODEL_NAME", "ibm-research/PowerMoE-3b") # Number of data parallel ranks for external LB testing DP_SIZE = int(os.getenv("DP_SIZE", "2")) diff --git a/tests/v1/e2e/general/test_async_scheduling.py b/tests/v1/e2e/general/test_async_scheduling.py index 28a1bedbe0b2..c3c4970de382 100644 --- a/tests/v1/e2e/general/test_async_scheduling.py +++ b/tests/v1/e2e/general/test_async_scheduling.py @@ -57,6 +57,8 @@ def test_without_spec_decoding( dict(bad_words=["the", " the"]), dict(logprobs=2), dict(logprobs=2, frequency_penalty=-1.0), + dict(prompt_logprobs=2), + dict(prompt_logprobs=2, logprobs=2), dict(structured_outputs=struct_outputs), dict( structured_outputs=struct_outputs, @@ -126,6 +128,8 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke dict(bad_words=["the", " the"]), dict(logprobs=2), dict(logprobs=2, frequency_penalty=-1.0), + dict(prompt_logprobs=2), + dict(prompt_logprobs=2, logprobs=2), dict(structured_outputs=struct_outputs), dict( structured_outputs=struct_outputs, @@ -413,7 +417,12 @@ def _all_logprobs_match(req_a, req_b) -> bool: ) -def _logprobs_match(lps_a: dict[int, Logprob], lps_b: dict[int, Logprob]) -> bool: +def _logprobs_match( + lps_a: dict[int, Logprob] | None, + lps_b: dict[int, Logprob] | None, +) -> bool: + if lps_a is None or lps_b is None: + return lps_a is lps_b rel_tol, abs_tol = 1e-3, 1e-6 return ( len(lps_a) == len(lps_b) diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 28fb2931b229..460e0d685649 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -33,11 +33,10 @@ SAMPLE_PROMPT = BatchLogprobsComposition.SAMPLE_PROMPT # On ROCm, floating-point reductions in attention and GEMM kernels are -# non-associative and sensitive to batch geometry. The ref LLM (no spec -# decode, default scheduling) and the spec-decode LLM (chunked prefill, -# different effective batch sizes) follow different reduction orders, -# producing numerically divergent logprobs that get misattributed to -# spec-decode incorrectness. +# non-associative and sensitive to batch geometry. If the ref LLM and +# spec-decode LLM use different scheduling or batch geometry, they can +# follow different reduction orders and produce numerically divergent +# logprobs that get misattributed to spec-decode incorrectness. # # Force LLM instances into an identical, deterministic execution # mode so the test isolates spec-decode correctness only: @@ -1086,18 +1085,25 @@ def test_spec_decode_logprobs( ) max_model_len = 256 - - # Run base LLM. - ref_llm = LLM( - model=model_name, + llm_kwargs = dict( max_logprobs=5, max_model_len=max_model_len, seed=42, logprobs_mode=logprobs_mode, gpu_memory_utilization=0.4, + # Force the same prefill chunking for both the base model and + # spec decode model so the comparison isolates spec decode. + enable_chunked_prefill=True, + max_num_batched_tokens=32, enable_prefix_caching=False, **ROCM_DETERMINISM_KWARGS, ) + + # Run base LLM. + ref_llm = LLM( + model=model_name, + **llm_kwargs, + ) ref_results = ref_llm.generate( [prompt, prompt], [sampling_params, penalty_sampling_params] ) @@ -1117,16 +1123,7 @@ def test_spec_decode_logprobs( spec_llm = LLM( model_name, speculative_config=spec_config_with_len, - max_logprobs=5, - max_model_len=max_model_len, - seed=42, - logprobs_mode=logprobs_mode, - gpu_memory_utilization=0.4, - # Force prefill chunking - enable_chunked_prefill=True, - max_num_batched_tokens=32, - enable_prefix_caching=False, - **ROCM_DETERMINISM_KWARGS, + **llm_kwargs, ) spec_results = spec_llm.generate( [prompt, prompt], [sampling_params, penalty_sampling_params] diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index 73ef8b915821..c0503fd69712 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -810,6 +810,9 @@ def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None "compute_capability": compute_cap, "is_mla": is_mla_backend or check_method_overrides(class_node, "is_mla"), "supports_sink": check_method_overrides(class_node, "supports_sink"), + "supports_non_causal": check_method_overrides( + class_node, "supports_non_causal" + ), "is_sparse": check_method_overrides(class_node, "is_sparse"), "supports_mm_prefix": check_method_overrides(class_node, "supports_mm_prefix"), "supports_dcp": supports_dcp, @@ -1311,6 +1314,10 @@ def _extract_priorities(body: list, priorities: dict[str, list[str]], prefix: st _COL_BLOCK_SIZES: TableColumn = ("Block Sizes", lambda b: b["block_sizes"]) _COL_HEAD_SIZES: TableColumn = ("Head Sizes", lambda b: b["head_sizes"]) _COL_SINK: TableColumn = ("Sink", lambda b: bool_to_emoji(b["supports_sink"])) +_COL_NON_CAUSAL: TableColumn = ( + "Non-Causal", + lambda b: bool_to_emoji(b["supports_non_causal"]), +) _COL_SPARSE: TableColumn = ("Sparse", lambda b: bool_to_emoji(b["is_sparse"])) _COL_MM_PREFIX: TableColumn = ( "MM Prefix", @@ -1344,6 +1351,7 @@ def _build_columns(is_mla: bool, has_versions: bool) -> list[TableColumn]: cols.append(_COL_VERSION) cols.extend([_COL_DTYPES, _COL_KV_DTYPES, _COL_BLOCK_SIZES, _COL_HEAD_SIZES]) cols.append(_COL_SINK) + cols.append(_COL_NON_CAUSAL) if is_mla: cols.append(_COL_SPARSE) cols.extend([_COL_MM_PREFIX, _COL_DCP, _COL_ATTN_TYPES, _COL_COMPUTE_CAP]) @@ -1554,6 +1562,7 @@ def generate_legend() -> str: | **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) | | **Head Sizes** | Supported attention head sizes | | **Sink** | Attention sink support (for StreamingLLM) | +| **Non-Causal** | Non-causal (bidirectional) attention support for decoder models | | **Sparse** | Sparse attention support (MLA only) | | **MM Prefix** | Multimodal prefix full attention support | | **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) | diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index 419275d2e6ae..b032c0a0d613 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -1803,7 +1803,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: if args.dataset_name == "custom": dataset = CustomDataset( - dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle + dataset_path=args.dataset_path, + disable_shuffle=args.disable_shuffle, + random_seed=args.seed, ) input_requests = dataset.sample( num_requests=args.num_prompts, @@ -1816,7 +1818,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: elif args.dataset_name == "custom_mm": dataset = CustomMMDataset( - dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle + dataset_path=args.dataset_path, + disable_shuffle=args.disable_shuffle, + random_seed=args.seed, ) input_requests = dataset.sample( num_requests=args.num_prompts, diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 6ba392802e31..95fd8787afe5 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -135,8 +135,10 @@ class ParallelConfig: data_parallel_external_lb: bool = False """Whether to use "external" DP LB mode. Applies only to online serving and when data_parallel_size > 0. This is useful for a "one-pod-per-rank" - wide-EP setup in Kubernetes. Set implicitly when --data-parallel-rank - is provided explicitly to vllm serve.""" + wide-EP setup in Kubernetes. Supported only for MoE deployments; non-MoE + models should use independent vLLM instances without --data-parallel-* + arguments. Set implicitly when --data-parallel-rank is provided explicitly + to vllm serve.""" data_parallel_hybrid_lb: bool = False """Whether to use "hybrid" DP LB mode. Applies only to online serving and when data_parallel_size > 0. Enables running an AsyncLLM diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 0146ee4c144a..52c04509b2fa 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -209,7 +209,9 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: "use_inductor_graph_partition": False, }, "kernel_config": { - "enable_flashinfer_autotune": True, + # Disabled for now due to correctness issues: + # https://github.com/flashinfer-ai/flashinfer/issues/3197 + "enable_flashinfer_autotune": False, }, } OPTIMIZATION_LEVEL_02 = { @@ -229,7 +231,9 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: "use_inductor_graph_partition": False, }, "kernel_config": { - "enable_flashinfer_autotune": True, + # Disabled for now due to correctness issues: + # https://github.com/flashinfer-ai/flashinfer/issues/3197 + "enable_flashinfer_autotune": False, }, } OPTIMIZATION_LEVEL_03 = { @@ -1613,7 +1617,7 @@ def _set_compile_ranges(self): max_size = rocm_aiter_ops.get_aiter_allreduce_max_size() else: max_size = compilation_config.pass_config.flashinfer_max_size(tp_size) - if max_size is not None: + if max_size is not None and self.model_config is not None: assert isinstance(self.model_config.dtype, torch.dtype) max_token_num = max_size // ( self.model_config.get_hidden_size() diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index cd9551003339..1b3803139217 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -962,7 +962,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "-dpn", type=int, help="Data parallel rank of this instance. " - "When set, enables external load balancer mode.", + "When set, enables external load balancer mode for MoE " + "data-parallel deployments. Unsupported for non-MoE models; " + "launch independent vLLM instances instead.", ) parallel_group.add_argument( "--data-parallel-start-rank", @@ -1697,29 +1699,15 @@ def create_engine_config( kv_offloading_backend=self.kv_offloading_backend, ) - # TurboQuant: auto-skip first/last 2 layers (boundary protection). - # These layers are most sensitive to quantization error. - # Users can add extra layers via --kv-cache-dtype-skip-layers. if resolved_cache_dtype.startswith("turboquant_"): - if model_config.is_hybrid: - raise NotImplementedError( - "TurboQuant KV cache is not supported for hybrid " - "(attention + Mamba) models. Boundary layer protection " - "requires uniform attention layers." - ) from vllm.model_executor.layers.quantization.turboquant.config import ( TurboQuantConfig, ) - num_layers = model_config.hf_text_config.num_hidden_layers - boundary = TurboQuantConfig.get_boundary_skip_layers(num_layers) + boundary = TurboQuantConfig.get_boundary_skip_layers(model_config) existing = set(cache_config.kv_cache_dtype_skip_layers) - merged = sorted(existing | set(boundary), key=lambda x: int(x)) - cache_config.kv_cache_dtype_skip_layers = merged - logger.info( - "TQ: skipping layers %s for boundary protection (num_layers=%d)", - merged, - num_layers, + cache_config.kv_cache_dtype_skip_layers = sorted( + existing | set(boundary), key=int ) ray_runtime_env = None @@ -1793,6 +1781,16 @@ def create_engine_config( data_parallel_external_lb = ( self.data_parallel_external_lb or self.data_parallel_rank is not None ) + if ( + self.data_parallel_size > 1 + and data_parallel_external_lb + and not model_config.is_moe + ): + raise ValueError( + "Non-MoE models do not support external data parallel mode. " + "For external load balancing, launch independent vLLM " + "instances without --data-parallel-* arguments." + ) # Local DP rank = 1, use pure-external LB. if data_parallel_external_lb: assert self.data_parallel_rank is not None, ( diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 9aac19e2fda5..da2ec10284c5 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -321,6 +321,21 @@ async def init_app_state( supported_tasks: tuple["SupportedTask", ...] | None = None, ) -> None: vllm_config = engine_client.vllm_config + + # Propagate enable_in_reasoning to the API-server process. The engine core + # runs in a separate process, so the contextvar that backs + # `get_current_vllm_config_or_none()` is None on this stack. Tool parsers + # call `get_enable_structured_outputs_in_reasoning()` during request + # handling and need to see the real flag, otherwise they silently fall + # back to False and mismatch the engine-side bitmask gating. + from vllm.tool_parsers.structural_tag_registry import ( + set_enable_structured_outputs_in_reasoning, + ) + + set_enable_structured_outputs_in_reasoning( + vllm_config.structured_outputs_config.enable_in_reasoning + ) + if supported_tasks is None: warnings.warn( "The 'supported_tasks' parameter was not provided to " diff --git a/vllm/envs.py b/vllm/envs.py index b2db5a8112bc..e456bec5bfb1 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -226,6 +226,7 @@ VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False VLLM_SYSTEM_START_DATE: str | None = None VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY: bool = False + VLLM_ENFORCE_STRICT_TOOL_CALLING: bool = False VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False VLLM_NVTX_SCOPES_FOR_PROFILING: bool = False VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES: bool = True @@ -1593,6 +1594,12 @@ def _get_or_set_default() -> str: "VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY": lambda: bool( int(os.getenv("VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY", "0")) ), + # When 1,the model structural tags will be used to enforce the model + # output conforming to the model's tool-calling format and schema. + # Default 0 (off). + "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: bool( + int(os.getenv("VLLM_ENFORCE_STRICT_TOOL_CALLING", "0")) + ), # Add optional custom scopes for profiling, disable to avoid overheads "VLLM_CUSTOM_SCOPES_FOR_PROFILING": lambda: bool( int(os.getenv("VLLM_CUSTOM_SCOPES_FOR_PROFILING", "0")) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py index eb14f9ec378c..81b1af85da5a 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py @@ -2,10 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import dataclasses +from fractions import Fraction +from typing import cast + import torch from vllm import _custom_ops as ops +from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + GPTQ_MARLIN_MIN_THREAD_N, MARLIN_SUPPORTED_GROUP_SIZES, apply_gptq_marlin_linear, check_marlin_supports_shape, @@ -23,11 +29,85 @@ from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ from vllm.platforms import current_platform from vllm.scalar_type import scalar_types +from vllm.utils.math_utils import round_up from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig +logger = init_logger(__name__) + + +def _pad_tensor_dim(x: torch.Tensor, dim: int, pad: int) -> torch.Tensor: + if pad == 0: + return x + + dim = dim if dim >= 0 else x.dim() + dim + pad_shape = list(x.shape) + pad_shape[dim] = pad + return torch.cat([x, x.new_zeros(pad_shape)], dim=dim) + + +def _get_param_output_dim_padding( + param: BasevLLMParameter, + output_dim_pad: int, +) -> tuple[int, int]: + output_dim = getattr(param, "output_dim", None) + if output_dim is None: + raise ValueError( + "Marlin output-dim padding requires vLLM parameter output_dim metadata." + ) + + pad = output_dim_pad + if getattr(param, "packed_dim", None) == output_dim: + packed_factor = getattr(param, "packed_factor", None) + if packed_factor is None: + raise ValueError( + "Marlin packed output-dim padding requires packed_factor metadata." + ) + + packed_pad = Fraction(output_dim_pad, 1) / Fraction(packed_factor) + if packed_pad.denominator != 1: + raise ValueError( + "Marlin output padding is not divisible by packed_factor: " + f"pad={output_dim_pad}, packed_factor={packed_factor}." + ) + pad = packed_pad.numerator + + return output_dim, pad + + +def _pad_parameter_output_dim( + param: BasevLLMParameter, + output_dim_pad: int, +) -> None: + output_dim, pad = _get_param_output_dim_padding(param, output_dim_pad) + param.data = _pad_tensor_dim(param.data, output_dim, pad) + class MarlinLinearKernel(MPLinearKernel): + config: MPLinearLayerConfig + w_q_name: str + w_s_name: str + w_zp_name: str | None + w_gidx_name: str | None + orig_output_size_per_partition: int + + def __init__( + self, + c: MPLinearLayerConfig, + w_q_param_name: str, + w_s_param_name: str, + w_zp_param_name: str | None = None, + w_gidx_param_name: str | None = None, + ) -> None: + super().__init__( + c, + w_q_param_name, + w_s_param_name, + w_zp_param_name, + w_gidx_param_name, + ) + self.orig_output_size_per_partition = self.config.partition_weight_shape[1] + @classmethod def get_min_capability(cls) -> int: return 75 @@ -54,17 +134,54 @@ def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: f"{MARLIN_SUPPORTED_GROUP_SIZES}", ) + # Pad sub-tile output dims at load time; see _maybe_pad_n(). + padded_n = round_up(c.partition_weight_shape[1], GPTQ_MARLIN_MIN_THREAD_N) return check_marlin_supports_shape( - c.partition_weight_shape[1], # out_features + padded_n, # out_features (possibly padded up to tile multiple) c.partition_weight_shape[0], # in_features c.full_weight_shape[0], # in_features c.group_size, ) + def _maybe_pad_n(self, layer: torch.nn.Module) -> None: + """Pad output dim to a Marlin tile multiple when needed.""" + c = self.config + orig_n = c.partition_weight_shape[1] + padded_n = round_up(orig_n, GPTQ_MARLIN_MIN_THREAD_N) + self.orig_output_size_per_partition = orig_n + if padded_n == orig_n: + return + + pad = padded_n - orig_n + + q = cast(BasevLLMParameter, getattr(layer, self.w_q_name)) + _pad_parameter_output_dim(q, pad) + + s = cast(BasevLLMParameter, getattr(layer, self.w_s_name)) + _pad_parameter_output_dim(s, pad) + + if c.zero_points and self.w_zp_name is not None: + zp = getattr(layer, self.w_zp_name, None) + if zp is not None: + _pad_parameter_output_dim(cast(BasevLLMParameter, zp), pad) + + self.config = dataclasses.replace( + c, + partition_weight_shape=(c.partition_weight_shape[0], padded_n), + ) + logger.info_once( + "Marlin: padded output dim %d -> %d to satisfy tile_n=%d", + orig_n, + padded_n, + GPTQ_MARLIN_MIN_THREAD_N, + ) + # note assumes that # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0} # `weight_scale` is: {input_dim = 0, output_dim = 1} def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self._maybe_pad_n(layer) + device = getattr(layer, self.w_q_name).device c = self.config is_a_8bit = c.act_type is not None and c.act_type.itemsize == 1 @@ -167,7 +284,10 @@ def transform_w_s(x): self._transform_param(layer, self.w_q_name, transform_w_q) self._transform_param(layer, self.w_s_name, transform_w_s) - if hasattr(layer, "bias") and layer.bias is not None: + is_output_padded = ( + self.orig_output_size_per_partition != self.config.partition_weight_shape[1] + ) + if hasattr(layer, "bias") and layer.bias is not None and not is_output_padded: layer.bias.data = marlin_permute_bias(layer.bias) def apply_weights( @@ -182,7 +302,22 @@ def apply_weights( # `process_weights_after_loading` will ensure w_zp and w_gidx are not # None for marlin - return apply_gptq_marlin_linear( + padded_n = c.partition_weight_shape[1] + orig_n = self.orig_output_size_per_partition + + if bias is not None and orig_n != padded_n: + if bias.shape[-1] == orig_n: + bias = _pad_tensor_dim(bias, -1, padded_n - orig_n) + bias = marlin_permute_bias(bias) + elif bias.shape[-1] == padded_n: + bias = marlin_permute_bias(bias) + else: + raise ValueError( + "Marlin bias shape does not match original or padded output dim: " + f"bias={bias.shape[-1]}, orig_n={orig_n}, padded_n={padded_n}." + ) + + out = apply_gptq_marlin_linear( input=x, weight=w_q, weight_scale=w_s, @@ -192,9 +327,13 @@ def apply_weights( workspace=self.workspace, wtype=c.weight_type, input_size_per_partition=c.partition_weight_shape[0], - output_size_per_partition=c.partition_weight_shape[1], + output_size_per_partition=padded_n, is_k_full=self.is_k_full, input_global_scale=getattr(layer, "input_global_scale", None), bias=bias, input_dtype=c.act_type, ) + + if orig_n != padded_n: + out = out[..., :orig_n].contiguous() + return out diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 82eecc8cd49b..20981f60cd24 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -238,6 +238,9 @@ kFp8StaticTensorSym, kNvfp4Dynamic, ) +from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import ( + yarn_get_mscale, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer from vllm.utils.math_utils import cdiv, round_down @@ -1327,6 +1330,35 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims: ) +def get_mla_prefill_scale(model_config: ModelConfig) -> float: + hf_text_config = model_config.hf_text_config + mla_dims = get_mla_dims(model_config) + qk_head_dim = mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim + scale = qk_head_dim**-0.5 + + # Deepseek V4 disables YaRN mscale for attention; Deepseek V2/V3 applies + # the same mscale correction when constructing the MLA attention module. + if hasattr(hf_text_config, "compress_ratios"): + return scale + + rope_parameters = getattr(hf_text_config, "rope_parameters", None) + if rope_parameters is None: + rope_parameters = getattr(hf_text_config, "rope_scaling", None) + + if rope_parameters is None: + return scale + + rope_type = rope_parameters.get("rope_type", rope_parameters.get("type")) + apply_yarn_scaling = rope_parameters.get("apply_yarn_scaling", True) + if rope_type != "default" and apply_yarn_scaling: + mscale_all_dim = rope_parameters.get("mscale_all_dim", False) + scaling_factor = rope_parameters["factor"] + mscale = yarn_get_mscale(float(scaling_factor), float(mscale_all_dim)) + scale *= mscale * mscale + + return scale + + @functools.cache def backend_supports_prefill_query_quantization() -> bool: """Check if the selected MLA prefill backend supports query quantization. @@ -1527,7 +1559,7 @@ def __init__( prefill_backend_cls = get_mla_prefill_backend(vllm_config) self._prefill_backend = prefill_backend_cls( num_heads=self.num_heads, - scale=self.model_config.get_head_size() ** -0.5, + scale=get_mla_prefill_scale(self.model_config), kv_lora_rank=self.mla_dims.kv_lora_rank, qk_nope_head_dim=self.mla_dims.qk_nope_head_dim, qk_rope_head_dim=self.mla_dims.qk_rope_head_dim, diff --git a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py index e10be4af8680..d6bd2b140087 100644 --- a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py @@ -62,7 +62,7 @@ def _supports_current_device() -> bool: @staticmethod def _supports_no_act_and_mul() -> bool: - return False + return True @staticmethod def _supports_activation(activation: MoEActivation) -> bool: @@ -70,6 +70,7 @@ def _supports_activation(activation: MoEActivation) -> bool: MoEActivation.SILU, MoEActivation.GELU, MoEActivation.SWIGLUOAI, + MoEActivation.RELU2_NO_MUL, ] @staticmethod diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 3de05cd93d36..456f40bbf7a3 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -538,9 +538,11 @@ def _get_quant_method() -> FusedMoEMethodBase: # for heuristic purposes, so it must be initialized first. self.quant_method: FusedMoEMethodBase = _get_quant_method() - if not self.moe_config.is_act_and_mul and not current_platform.is_cuda_alike(): + if not self.moe_config.is_act_and_mul and not ( + current_platform.is_cuda_alike() or current_platform.is_xpu() + ): raise NotImplementedError( - "is_act_and_mul=False is supported only for CUDA and ROCm for now" + "is_act_and_mul=False is supported only for CUDA and XPU for now" ) if self.enable_eplb and not self.quant_method.supports_eplb: diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 01ac5cfa9da7..f4796243e013 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -381,7 +381,7 @@ def convert_to_nvfp4_moe_kernel_format( elif nvfp4_backend == NvFp4MoeBackend.EMULATION: # Move the E2M1 lookup table to the device now, because # `.to(device)` is not allowed during CUDA graph capture. - kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(layer.weight.device) + kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(w13.device) if a13_scale is None or a2_scale is None: raise ValueError( diff --git a/vllm/model_executor/layers/quantization/turboquant/config.py b/vllm/model_executor/layers/quantization/turboquant/config.py index f9cfc89c0c1d..50beb8d1d9bf 100644 --- a/vllm/model_executor/layers/quantization/turboquant/config.py +++ b/vllm/model_executor/layers/quantization/turboquant/config.py @@ -2,8 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """TurboQuant configuration.""" +from __future__ import annotations + +import logging import math from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from vllm.config import ModelConfig + +logger = logging.getLogger(__name__) # Named TQ presets: each maps to frozen config parameters. # key_quant_bits: 8 = FP8 keys, 3-4 = MSE (Lloyd-Max) quantized keys. @@ -159,12 +168,34 @@ def slot_size_aligned(self) -> int: return s + (s % 2) # round up to even @staticmethod - def get_boundary_skip_layers(num_layers: int, n: int = 2) -> list[str]: - """Get layer indices to skip TQ compression (boundary protection). - - Returns first N and last N layer indices as strings, suitable for - kv_cache_dtype_skip_layers. + def get_boundary_skip_layers( + model_config: ModelConfig, + n: int = 2, + ) -> list[str]: + """Layer indices to skip TQ compression (boundary protection). + + For hybrid models (attention + Mamba/linear-attention), boundary + protection is disabled — hybrids typically have only 8-12 + full-attention layers and a hard n=2 on each side would cover + ~40 % of them. The dense GSM8K baselines that motivate n=2 + don't apply to hybrids. + + For dense models, skips first N and last N attention layers. + Empirically required for aggressive presets (k3v4_nc, 3bit_nc) + — without it GSM8K drops ~30 points on Qwen3-4B. """ + if model_config.is_hybrid: + attn_indices = _get_full_attention_layer_indices(model_config) + if not attn_indices: + raise NotImplementedError( + "TurboQuant KV cache requires identifiable " + "full-attention layers, but none were found in " + "the hybrid model config." + ) + logger.info("TQ hybrid: full-attention layers %s", attn_indices) + return [] + + num_layers = model_config.hf_text_config.num_hidden_layers if n <= 0 or num_layers <= 0: return [] n = min(n, num_layers // 2) # don't skip more than half @@ -175,7 +206,7 @@ def get_boundary_skip_layers(num_layers: int, n: int = 2) -> list[str]: return [str(i) for i in indices] @staticmethod - def from_cache_dtype(cache_dtype: str, head_dim: int) -> "TurboQuantConfig": + def from_cache_dtype(cache_dtype: str, head_dim: int) -> TurboQuantConfig: """Create config from a named preset. Valid presets: turboquant_k8v4, turboquant_4bit_nc, etc. @@ -193,3 +224,31 @@ def from_cache_dtype(cache_dtype: str, head_dim: int) -> "TurboQuantConfig": value_quant_bits=preset["value_quant_bits"], norm_correction=preset["norm_correction"], ) + + +def _get_full_attention_layer_indices(model_config: ModelConfig) -> list[int]: + """Global indices of full-attention layers in a hybrid model. + + Covers the conventions used across vLLM: ``layer_types`` (Qwen3.5/Next), + ``layers_block_type`` (Jamba/Zamba2), ``attn_type_list`` (Minimax). + """ + text_cfg = model_config.hf_text_config + hf_cfg = model_config.hf_config + + layer_types = getattr(text_cfg, "layer_types", None) + if layer_types is not None: + return [ + i for i, t in enumerate(layer_types) if t in ("full_attention", "attention") + ] + + layers_block_type = getattr(text_cfg, "layers_block_type", None) + if layers_block_type is not None: + return [ + i for i, t in enumerate(layers_block_type) if t in ("attention", "hybrid") + ] + + attn_type_list = getattr(hf_cfg, "attn_type_list", None) + if attn_type_list is not None: + return [i for i, t in enumerate(attn_type_list) if t == 1] + + return [] diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index d659effd70ff..765b93dba55a 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.platforms import current_platform from vllm.scalar_type import ScalarType, scalar_types +from vllm.utils.math_utils import round_up from vllm.utils.platform_utils import num_compute_units from .quant_utils import pack_cols, unpack_cols @@ -218,8 +219,16 @@ def check_marlin_supports_layer(layer: LinearBase, group_size: int) -> bool: getattr(layer, "input_size_per_partition", None) or layer.input_size ) + # MarlinLinearKernel pads sub-tile output partitions at load time. + # Keep this pre-selection helper aligned with can_implement(), otherwise + # AutoRound/INC and AWQMarlin can incorrectly reject small-N layers that + # the Marlin kernel handles by padding and slicing back to the original N. + padded_output_size_per_partition = round_up( + output_size_per_partition, GPTQ_MARLIN_MIN_THREAD_N + ) + return check_marlin_supports_shape( - output_size_per_partition=output_size_per_partition, + output_size_per_partition=padded_output_size_per_partition, input_size_per_partition=input_size_per_partition, input_size=layer.input_size, group_size=group_size, diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index 91fce6f57b3e..baa2081d58b2 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/vllm/model_executor/model_loader/reload/meta.py @@ -102,7 +102,7 @@ def materialize_layer(layer: torch.nn.Module, info: LayerReloadingInfo): with info.restore_device: for name, tensor in get_layer_tensors(layer).items(): - if name not in SKIP_TENSORS: + if name not in SKIP_TENSORS and tensor.is_meta: setattr(layer, name, materialize_meta_tensor(tensor)) diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 684ced0a6abd..994b52606b18 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -1499,6 +1499,11 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + mm_config = self.model_config.multimodal_config + load_multimodal_weights = not all( + mm_config.get_limit_per_prompt(modality) == 0 + for modality in ("image", "video", "audio") + ) adapter_dict = dict(self.mlp1.named_parameters()) def is_llm(name: str) -> bool: @@ -1523,23 +1528,30 @@ def is_sound_weights(name: str) -> bool: # Strip 'language_model.' prefix for LLM weights llm_weights.append((".".join(name.split(".")[1:]), w)) elif is_adapter_weights((name, w)): + if not load_multimodal_weights: + continue # Load vision-language adapter weights directly trimmed_name = ".".join(name.split(".")[1:]) param = adapter_dict[trimmed_name] with torch.no_grad(): default_weight_loader(param, w) elif is_vision_weights(name): + if not load_multimodal_weights: + continue # Convert: vision_model.radio_model.* → radio_model.* hf_key = name[len("vision_model.") :] # Remove "vision_model." prefix vision_weights.append((hf_key, w)) elif is_sound_weights(name): + if not load_multimodal_weights: + continue assert self.sound_encoder is not None sound_weights.append((name, w)) self.language_model.load_weights(llm_weights) - self.vision_model.load_weights(vision_weights) - if self.sound_encoder is not None and len(sound_weights) > 0: - self.sound_encoder.load_weights(sound_weights) + if load_multimodal_weights: + self.vision_model.load_weights(vision_weights) + if self.sound_encoder is not None and len(sound_weights) > 0: + self.sound_encoder.load_weights(sound_weights) def get_vit_model_from_radio_config(self, hf_config): hf_config_vision = hf_config.vision_config diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 2753326755fb..80952ced73d1 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -545,6 +545,42 @@ def _align_hybrid_block_size( dtype=kv_cache_dtype, kv_quant_mode=kv_quant_mode, ).page_size_bytes + elif cache_config.cache_dtype.startswith("turboquant_"): + # TQ has a packed K|V layout; the standard FullAttentionSpec + # formula over-sizes it and trips unify_kv_cache_spec_page_size + # when all attention layers are TQ. With mixed skip+TQ the skip + # layers still use the standard layout — take max so mamba + # padding covers the largest actual page. + from vllm.model_executor.layers.quantization.turboquant.config import ( + TurboQuantConfig, + ) + from vllm.v1.kv_cache_interface import TQFullAttentionSpec + + tq_cfg = TurboQuantConfig.from_cache_dtype( + cache_config.cache_dtype, model_config.get_head_size() + ) + tq_page = TQFullAttentionSpec( + block_size=1, + num_kv_heads=model_config.get_num_kv_heads(parallel_config), + head_size=model_config.get_head_size(), + head_size_v=model_config.get_head_size(), + dtype=kv_cache_dtype, + kv_quant_mode=kv_quant_mode, + tq_slot_size=tq_cfg.slot_size_aligned, + ).page_size_bytes + if cache_config.kv_cache_dtype_skip_layers: + skip_page = FullAttentionSpec( + block_size=1, + num_kv_heads=model_config.get_num_kv_heads(parallel_config), + head_size=model_config.get_head_size(), + dtype=model_config.dtype, + ).page_size_bytes + # lcm, not max: skip_page is often not a multiple of + # tq_page, so max would leave per-layer page sizes + # un-unifiable downstream. + attn_page_size_1_token = lcm(tq_page, skip_page) + else: + attn_page_size_1_token = tq_page else: attn_page_size_1_token = FullAttentionSpec( block_size=1, diff --git a/vllm/reasoning/kimi_k2_reasoning_parser.py b/vllm/reasoning/kimi_k2_reasoning_parser.py index 7a92703426fc..0b64c5c62ea1 100644 --- a/vllm/reasoning/kimi_k2_reasoning_parser.py +++ b/vllm/reasoning/kimi_k2_reasoning_parser.py @@ -221,6 +221,10 @@ def extract_reasoning_streaming( return None if self._end_token_id in delta_token_ids: + if self._end_token not in delta_text: + # Token ID arrived before text was flushed (stop-sequence buffering). + # Wait for the next delta when the text becomes visible. + return None end_index = delta_text.find(self._end_token) reasoning = delta_text[:end_index] content = delta_text[end_index + len(self._end_token) :] @@ -229,6 +233,9 @@ def extract_reasoning_streaming( ) if self._tool_section_start_token_id in delta_token_ids: + if self._tool_section_start_token not in delta_text: + # Token ID arrived before text was flushed (stop-sequence buffering). + return None tool_index = delta_text.find(self._tool_section_start_token) reasoning = delta_text[:tool_index] content = delta_text[tool_index:] diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index 75181d8dfac6..c3438082a72d 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/vllm/tool_parsers/abstract_tool_parser.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib +import json import os from collections.abc import Callable, Sequence from functools import cached_property @@ -13,6 +14,7 @@ from openai.types.responses.function_tool import FunctionTool from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionToolsParam, ) @@ -23,6 +25,7 @@ from vllm.entrypoints.openai.responses.protocol import ( ResponsesRequest, ) +from vllm.envs import VLLM_ENFORCE_STRICT_TOOL_CALLING from vllm.logger import init_logger from vllm.sampling_params import ( StructuredOutputsParams, @@ -83,13 +86,39 @@ def vocab(self) -> dict[str, int]: return self.model_tokenizer.get_vocab() def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest + self, + request: ChatCompletionRequest | ResponsesRequest, ) -> ChatCompletionRequest | ResponsesRequest: - """ - Static method that used to adjust the request parameters. - """ + # If there are no tools, return the request as is. if not request.tools: return request + + # Step 1 (highest priority for ChatCompletionRequest): apply + # vLLM-owned structural tag support for model-specific tool formats. + if ( + isinstance(request, ChatCompletionRequest) + and VLLM_ENFORCE_STRICT_TOOL_CALLING + ): + need_tool_calling = ( + request.tool_choice == "auto" + or request.tool_choice == "required" + or isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) + ) + if need_tool_calling: + structure_tag = self.get_structural_tag(request) + if structure_tag is not None: + if request.structured_outputs is None: + request.structured_outputs = StructuredOutputsParams( + structural_tag=json.dumps(structure_tag.model_dump()), + ) + else: + request.structured_outputs.structural_tag = json.dumps( + structure_tag.model_dump() + ) + return request + + # Step 2: set structured output params when tool constraints are + # derived from the tool schema. json_schema_from_tool = get_json_schema_from_tools( tool_choice=request.tool_choice, tools=request.tools ) @@ -121,6 +150,9 @@ def adjust_request( return request + def get_structural_tag(self, request: ChatCompletionRequest): + return None + def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest ) -> ExtractedToolCallInformation: diff --git a/vllm/tool_parsers/deepseekv4_tool_parser.py b/vllm/tool_parsers/deepseekv4_tool_parser.py index 45a9c1302578..e32451cd8bbd 100644 --- a/vllm/tool_parsers/deepseekv4_tool_parser.py +++ b/vllm/tool_parsers/deepseekv4_tool_parser.py @@ -1,7 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) from vllm.tool_parsers.deepseekv32_tool_parser import DeepSeekV32ToolParser +from vllm.tool_parsers.structural_tag_registry import ( + get_enable_structured_outputs_in_reasoning, + get_model_structural_tag, +) class DeepSeekV4ToolParser(DeepSeekV32ToolParser): @@ -14,3 +21,11 @@ class DeepSeekV4ToolParser(DeepSeekV32ToolParser): tool_call_start_token: str = "<｜DSML｜tool_calls>" tool_call_end_token: str = "" + + def get_structural_tag(self, request: ChatCompletionRequest): + return get_model_structural_tag( + model="deepseek_v4", + tools=request.tools, + tool_choice=request.tool_choice, + reasoning=get_enable_structured_outputs_in_reasoning(), + ) diff --git a/vllm/tool_parsers/qwen3coder_tool_parser.py b/vllm/tool_parsers/qwen3coder_tool_parser.py index 7b089ceffbc0..73850b2ab0c5 100644 --- a/vllm/tool_parsers/qwen3coder_tool_parser.py +++ b/vllm/tool_parsers/qwen3coder_tool_parser.py @@ -25,12 +25,18 @@ Tool, ToolParser, ) +from vllm.tool_parsers.structural_tag_registry import ( + get_enable_structured_outputs_in_reasoning, + get_model_structural_tag, +) from vllm.tool_parsers.utils import find_tool_properties logger = init_logger(__name__) class Qwen3CoderToolParser(ToolParser): + supports_required_and_named: bool = False + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) @@ -681,3 +687,11 @@ def extract_tool_calls_streaming( return result return None + + def get_structural_tag(self, request: ChatCompletionRequest): + return get_model_structural_tag( + model="qwen_3_5", + tools=request.tools, + tool_choice=request.tool_choice, + reasoning=get_enable_structured_outputs_in_reasoning(), + ) diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py new file mode 100644 index 000000000000..754cc52361c5 --- /dev/null +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Model-specific structural tag builders adapted from XGrammar's +# builtin structural tag implementations: +# https://github.com/mlc-ai/xgrammar/blob/main/python/xgrammar/builtin_structural_tag.py + +from collections.abc import Callable +from typing import Any, Literal + +from xgrammar import StructuralTag +from xgrammar.structural_tag import ( + AnyTextFormat, + ConstStringFormat, + JSONSchemaFormat, + SequenceFormat, + TagFormat, + TagsWithSeparatorFormat, + TriggeredTagsFormat, +) + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, + ChatCompletionToolsParam, +) + +SimplifiedToolChoice = Literal["auto", "required", "forced"] +ToolChoice = ( + Literal["none", "auto", "required"] | ChatCompletionNamedToolChoiceParam | None +) +StructuralTagBuilder = Callable[ + [list[ChatCompletionToolsParam], SimplifiedToolChoice, bool], + StructuralTag, +] + +_structural_tag_registry: dict[str, StructuralTagBuilder] = {} + + +def register_model_structural_tag(name: str): + """Register a vLLM-owned model-specific structural tag builder.""" + + def decorator(func: StructuralTagBuilder) -> StructuralTagBuilder: + _structural_tag_registry[name] = func + return func + + return decorator + + +def get_model_structural_tag( + model: str, + tools: list[ChatCompletionToolsParam] | None, + tool_choice: ToolChoice, + reasoning: bool, +) -> StructuralTag | None: + """Build a structural tag from vLLM-owned model-specific builders.""" + + builder = _structural_tag_registry.get(model) + if builder is None: + supported = list(_structural_tag_registry.keys()) + raise ValueError(f"Unknown format type: {model}, supported types: {supported}") + + normalized_tools, simplified_tool_choice = _normalize_tool_choice( + tools=tools, + tool_choice=tool_choice, + ) + if not normalized_tools: + return None + + return builder(normalized_tools, simplified_tool_choice, reasoning) + + +def _normalize_tool_choice( + tools: list[ChatCompletionToolsParam] | None, + tool_choice: ToolChoice, +) -> tuple[list[ChatCompletionToolsParam], SimplifiedToolChoice]: + """Normalize vLLM ChatCompletion tool_choice for structural tag builders.""" + + if not tools: + return [], "auto" + + if tool_choice is None or tool_choice == "none": + return [], "auto" + + if tool_choice == "auto": + return tools, "auto" + + if tool_choice == "required": + return tools, "required" + + if isinstance(tool_choice, ChatCompletionNamedToolChoiceParam): + tool_name = tool_choice.function.name + filtered_tools = [tool for tool in tools if tool.function.name == tool_name] + if not filtered_tools: + raise ValueError( + f"The tool with name '{tool_name}' is not found in the tools list." + ) + return filtered_tools, "forced" + + raise ValueError(f"Unsupported tool_choice for structural tag: {tool_choice}") + + +def _get_function_parameters(function: Any) -> dict[str, Any] | bool: + """Return the JSON schema used for constrained tool arguments.""" + + if getattr(function, "strict", None) is False: + return True + if function.parameters is None: + return True + return function.parameters + + +_enable_structured_outputs_in_reasoning: bool = False + + +def set_enable_structured_outputs_in_reasoning(enabled: bool) -> None: + """Publish the engine's ``enable_in_reasoning`` flag to tool parsers. + + Called once during APIServer startup so request-time parsers can read + it without going through the EngineCore-only contextvar. + """ + + global _enable_structured_outputs_in_reasoning + _enable_structured_outputs_in_reasoning = bool(enabled) + + +def get_enable_structured_outputs_in_reasoning() -> bool: + """Whether structured outputs are active during the reasoning phase. + + When ``True``, the structural tag will cover the reasoning part: + ``...`` prefix (if available); when ``False`` (default), the tag only + constrains the post-reasoning suffix. + """ + + return _enable_structured_outputs_in_reasoning + + +@register_model_structural_tag("deepseek_v4") +def get_deepseek_v4_structural_tag( + tools: list[ChatCompletionToolsParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, +) -> StructuralTag: + """Build DeepSeek V4 structural tags.""" + + invoke_begin_prefix = '<｜DSML｜invoke name="' + invoke_begin_suffix = '">\n' + invoke_end = "\n" + tool_calls_prefix = "\n\n" + function_calls_begin = "<｜DSML｜tool_calls>\n" + function_calls_end = "" + function_calls_trigger = "<｜DSML｜tool_calls>" + think_tag_end = "" + think_exclude_tokens = ["", ""] + xml_style = "deepseek_xml" + + if tool_choice == "auto": + tags = [] + for tool in tools: + function = tool.function + parameters = _get_function_parameters(function) + tags.append( + TagFormat( + begin=invoke_begin_prefix + function.name + invoke_begin_suffix, + content=JSONSchemaFormat( + json_schema=parameters, + style=xml_style, + ), + end=invoke_end, + ) + ) + + if tags: + function_calling_tags = TagsWithSeparatorFormat( + tags=tags, + separator="\n", + at_least_one=True, + ) + suffix_tag = TriggeredTagsFormat( + triggers=[function_calls_trigger], + tags=[ + TagFormat( + begin=function_calls_begin, + content=function_calling_tags, + end=function_calls_end, + ) + ], + excludes=think_exclude_tokens, + ) + else: + suffix_tag = AnyTextFormat(excludes=think_exclude_tokens) + + elif tool_choice == "forced": + if not tools: + raise ValueError("Forced tool choice must resolve to exactly one tool.") + function = tools[0].function + suffix_tag = SequenceFormat( + elements=[ + ConstStringFormat(value=tool_calls_prefix + function_calls_begin), + TagFormat( + begin=invoke_begin_prefix + function.name + invoke_begin_suffix, + content=JSONSchemaFormat( + json_schema=_get_function_parameters(function), + style=xml_style, + ), + end=invoke_end, + ), + ConstStringFormat(value=function_calls_end), + ] + ) + + elif tool_choice == "required": + tags = [] + for tool in tools: + function = tool.function + parameters = _get_function_parameters(function) + tags.append( + TagFormat( + begin=invoke_begin_prefix + function.name + invoke_begin_suffix, + content=JSONSchemaFormat( + json_schema=parameters, + style=xml_style, + ), + end=invoke_end, + ) + ) + assert len(tags) > 0 + suffix_tag = SequenceFormat( + elements=[ + ConstStringFormat(value=tool_calls_prefix + function_calls_begin), + TagsWithSeparatorFormat( + tags=tags, + separator="\n", + at_least_one=True, + ), + ConstStringFormat(value=function_calls_end), + ] + ) + + if not reasoning: + return StructuralTag(format=suffix_tag) + + prefix_tag = TagFormat(begin="", content=AnyTextFormat(), end=think_tag_end) + return StructuralTag(format=SequenceFormat(elements=[prefix_tag, suffix_tag])) + + +@register_model_structural_tag("qwen_3_5") +def get_qwen_3_5_structural_tag( + tools: list[ChatCompletionToolsParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, +) -> StructuralTag: + """Build Qwen XML structural tags. + + This format is used for Qwen3-Coder/Qwen3.5/Qwen3.6 and is compatible with + Qwen variants that use the same XML tool-call format. + """ + tool_call_begin_prefix = "\n", ""] + + if tool_choice == "auto": + tags = [] + for tool in tools: + function = tool.function + parameters = _get_function_parameters(function) + tags.append( + TagFormat( + begin=f"{tool_call_begin_prefix}{function.name}{tool_call_begin_suffix}", + content=JSONSchemaFormat(json_schema=parameters, style="qwen_xml"), + end=tool_call_end, + ) + ) + + if tags: + suffix_tag = TriggeredTagsFormat( + triggers=[tool_call_trigger], + tags=tags, + excludes=think_exclude_tokens, + ) + else: + suffix_tag = AnyTextFormat(excludes=think_exclude_tokens) + + elif tool_choice == "forced": + if not tools: + raise ValueError("Forced tool choice must resolve to exactly one tool.") + function = tools[0].function + suffix_tag = TagFormat( + begin=f"{tool_call_begin_prefix}{function.name}{tool_call_begin_suffix}", + content=JSONSchemaFormat( + json_schema=_get_function_parameters(function), + style="qwen_xml", + ), + end=tool_call_end, + ) + + elif tool_choice == "required": + tags = [] + for tool in tools: + function = tool.function + parameters = _get_function_parameters(function) + tags.append( + TagFormat( + begin=f"{tool_call_begin_prefix}{function.name}{tool_call_begin_suffix}", + content=JSONSchemaFormat(json_schema=parameters, style="qwen_xml"), + end=tool_call_end, + ) + ) + assert len(tags) > 0 + suffix_tag = TagsWithSeparatorFormat( + tags=tags, + separator="", + at_least_one=True, + ) + + if not reasoning: + result = StructuralTag(format=suffix_tag) + else: + prefix_tag = SequenceFormat( + elements=[ + TagFormat(begin="", content=AnyTextFormat(), end=think_tag_end), + ConstStringFormat(value=think_suffix), + ] + ) + result = StructuralTag(format=SequenceFormat(elements=[prefix_tag, suffix_tag])) + + return result diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 2f00178ba6ef..c95df9c1077c 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -101,6 +101,7 @@ def __getitem__(self, key): fireredlid="FireRedLIDConfig", funaudiochat="FunAudioChatConfig", granite4_vision="Granite4VisionConfig", + hyperclovax_vlm="HCXVisionConfig", hunyuan_vl="HunYuanVLConfig", hy_v3="HYV3Config", isaac="IsaacConfig", @@ -217,8 +218,9 @@ def parse( ) else: if model_type in _CONFIG_REGISTRY: - # Register the config class to AutoConfig to ensure it's used in future - # calls to `from_pretrained` + # Register the config class to AutoConfig to ensure it's used + # in future calls to `from_pretrained` (e.g. from + # AutoTokenizer or AutoProcessor). config_class = _CONFIG_REGISTRY[model_type] config_class.model_type = model_type AutoConfig.register(model_type, config_class, exist_ok=True) diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 44abe32c916f..99f099adc786 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -37,6 +37,7 @@ "HunYuanVLConfig": "vllm.transformers_utils.configs.hunyuan_vl", "HunYuanVLTextConfig": "vllm.transformers_utils.configs.hunyuan_vl", "HunYuanVLVisionConfig": "vllm.transformers_utils.configs.hunyuan_vl", + "HCXVisionConfig": "vllm.transformers_utils.configs.hyperclovax", "HYV3Config": "vllm.transformers_utils.configs.hy_v3", "HyperCLOVAXConfig": "vllm.transformers_utils.configs.hyperclovax", "IsaacConfig": "vllm.transformers_utils.configs.isaac", @@ -104,6 +105,7 @@ "HunYuanVLConfig", "HunYuanVLTextConfig", "HunYuanVLVisionConfig", + "HCXVisionConfig", "HYV3Config", "HyperCLOVAXConfig", "IsaacConfig", diff --git a/vllm/transformers_utils/configs/hyperclovax.py b/vllm/transformers_utils/configs/hyperclovax.py index 9fa823743d66..d1a3218fe4dd 100644 --- a/vllm/transformers_utils/configs/hyperclovax.py +++ b/vllm/transformers_utils/configs/hyperclovax.py @@ -17,6 +17,7 @@ # limitations under the License. """HyperCLOVA X model configuration.""" +from transformers import AutoConfig from transformers.configuration_utils import PretrainedConfig @@ -275,3 +276,74 @@ def __init__( auto_map=auto_map, **kwargs, ) + + +class HCXVisionConfig(PretrainedConfig): + """Vendored HyperCLOVAX Vision config with transformers v5 fix. + + The original remote code config does not handle empty initialization + (text_config=None), which breaks transformers v5's @strict validation. + + TODO: Remove this class once HyperCLOVAX is upstreamed to transformers. + Tracking PR: https://github.com/huggingface/transformers/pull/44956 + """ + + model_type = "hyperclovax_vlm" + keys_to_ignore_at_inference = ["past_key_values"] + + text_config_attribute_map = { + "n_embd": "hidden_size", + "n_positions": "max_position_embeddings", + "n_head": "num_attention_heads", + "n_layer": "num_hidden_layers", + } + + def __init__( + self, + text_config=None, + vision_config=None, + use_nth_layer=-2, + img_start_id=100009, + decoder_max_length=4096, + anyres=False, + unpad=False, + max_num_grids=-1, + num_queries_vis_abstractor=-1, + ignore_index=-100, + proj_pos_emb=True, + proj_prenorm=False, + use_1x1_grid=False, + **kwargs, + ): + for key, val in self.text_config_attribute_map.items(): + if text_config is not None and key in text_config: + text_config[val] = text_config.pop(key) + + self.text_config = None + if text_config is not None: + _text_config = AutoConfig.for_model(text_config["model_type"]) + self.text_config = _text_config.from_dict(text_config) + self.hidden_size = self.text_config.hidden_size + + self.vision_config = None + if vision_config is not None: + _vision_config = AutoConfig.for_model(vision_config["model_type"]) + self.vision_config = _vision_config.from_dict(vision_config) + + self.use_nth_layer = use_nth_layer + self.decoder_max_length = decoder_max_length + self.anyres = anyres + self.unpad = unpad + self.max_num_grids = max_num_grids + self.num_queries_vis_abstractor = num_queries_vis_abstractor + self.img_start_id = img_start_id + self.ignore_index = ignore_index + self.proj_pos_emb = proj_pos_emb + self.proj_prenorm = proj_prenorm + self.use_1x1_grid = use_1x1_grid + super().__init__(**kwargs) + + def get_text_config(self, decoder=False): + if self.text_config is not None: + return self.text_config + return self diff --git a/vllm/utils/cpu_resource_utils.py b/vllm/utils/cpu_resource_utils.py index bbf554d0ccdd..25c299a0c0c1 100644 --- a/vllm/utils/cpu_resource_utils.py +++ b/vllm/utils/cpu_resource_utils.py @@ -125,7 +125,7 @@ def get_allowed_cpu_list() -> list[LogicalCPUInfo]: if platform.system() == "Darwin": return cpu_list - global_allowed_cpu_id_list = os.sched_getaffinity(0) + global_allowed_cpu_id_list = os.sched_getaffinity(0) # type: ignore[attr-defined] logical_cpu_list = [x for x in cpu_list if x.id in global_allowed_cpu_id_list] return logical_cpu_list diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 1eb9306ed4b1..798c136fc239 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -110,6 +110,32 @@ def is_strictly_contiguous(t: torch.Tensor) -> bool: return True +def canonicalize_singleton_dim_strides(t: torch.Tensor) -> torch.Tensor: + """Fix degenerate strides on size=1 dimensions for CUDA TMA compatibility. + + PyTorch allows any stride on a size=1 dim (is_contiguous() is always True + there), so a size=1 dim may have stride=1 (2 bytes for bf16) instead of + the canonical product(shape[i+1:]). CUDA TMA on H100+ requires all + non-outermost strides to be ≥16-byte aligned; stride=1 triggers + cudaErrorIllegalInstruction. Zero-copy: patches stride metadata only via + as_strided; returns t unchanged if all size=1 strides are already canonical. + """ + if 1 not in t.shape: + return t + strides = list(t.stride()) + shape = t.shape + prev_stride = 1 + changed = False + for i in range(len(shape) - 1, -1, -1): + if shape[i] == 1 and strides[i] != prev_stride: + strides[i] = prev_stride + changed = True + prev_stride = strides[i] * shape[i] + if not changed: + return t + return t.as_strided(t.shape, strides) + + @contextlib.contextmanager def set_default_torch_dtype(dtype: torch.dtype): """Sets the default torch dtype to the given dtype.""" diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 1c9ff3f79e43..e73954ee7478 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -11,7 +11,10 @@ from vllm.model_executor.layers.attention import Attention from vllm.platforms import current_platform -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import ( + canonicalize_singleton_dim_strides, + is_quantized_kv_cache, +) from vllm.v1.attention.backend import ( AttentionBackend, AttentionImpl, @@ -747,6 +750,23 @@ def forward( # For decoder and cross-attention, use KV cache as before key_cache, value_cache = kv_cache.unbind(0) + # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP). + # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment. + # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides. + fixed_k = canonicalize_singleton_dim_strides(key_cache) + fixed_v = canonicalize_singleton_dim_strides(value_cache) + if fixed_k is not key_cache or fixed_v is not value_cache: + logger.debug( + "Canonicalized degenerate KV cache strides (FlashAttention): " + "shape=%s, key strides before=%s after=%s, " + "value strides before=%s after=%s", + key_cache.shape, + key_cache.stride(), + fixed_k.stride(), + value_cache.stride(), + fixed_v.stride(), + ) + key_cache, value_cache = fixed_k, fixed_v if is_quantized_kv_cache(self.kv_cache_dtype): # queries are quantized in the attention layer @@ -861,6 +881,8 @@ def do_kv_cache_update( # we use direct Q, K, V tensors without caching return + # Scatter write into the KV cache using slot_mapping indices. + # No TMA kernel is invoked here, so stride canonicalization is not needed. key_cache, value_cache = kv_cache.unbind(0) # Reshape the input keys and values and store them in the cache. diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py index d18054769711..82a9f07a4e59 100644 --- a/vllm/v1/attention/backends/flash_attn_diffkv.py +++ b/vllm/v1/attention/backends/flash_attn_diffkv.py @@ -4,7 +4,11 @@ import torch -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.logger import init_logger +from vllm.utils.torch_utils import ( + canonicalize_singleton_dim_strides, + is_quantized_kv_cache, +) from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.fa_utils import ( get_flash_attn_version, @@ -25,6 +29,8 @@ cascade_attention, ) +logger = init_logger(__name__) + class FlashAttentionDiffKVBackend(FlashAttentionBackend): # Default to 128 for this backend @@ -204,6 +210,23 @@ def forward( # Different head_size for K and V key_cache = kv_cache[..., : self.head_size] value_cache = kv_cache[..., self.head_size :] + # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP). + # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment. + # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides. + fixed_k = canonicalize_singleton_dim_strides(key_cache) + fixed_v = canonicalize_singleton_dim_strides(value_cache) + if fixed_k is not key_cache or fixed_v is not value_cache: + logger.debug( + "Canonicalized degenerate KV cache strides (FlashAttentionDiffKV): " + "shape=%s, key strides before=%s after=%s, " + "value strides before=%s after=%s", + key_cache.shape, + key_cache.stride(), + fixed_k.stride(), + value_cache.stride(), + fixed_v.stride(), + ) + key_cache, value_cache = fixed_k, fixed_v if is_quantized_kv_cache(self.kv_cache_dtype): # queries are quantized in the attention layer diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 8f5cb6206bd0..2de61a2b1f28 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -43,6 +43,7 @@ from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.torch_utils import ( + canonicalize_singleton_dim_strides, is_quantized_kv_cache, is_strictly_contiguous, nvfp4_kv_cache_full_dim, @@ -1479,6 +1480,21 @@ def forward( stride_order = FlashInferBackend.get_kv_cache_stride_order() kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous + # Fix degenerate strides on any size-1 dimension (e.g. num_kv_heads=1 + # with TP=8). PyTorch permits non-canonical strides on size-1 dims; + # CUDA TMA requires ≥16-byte alignment on all non-outermost strides. + # canonicalize_singleton_dim_strides patches metadata via as_strided — + # zero-copy. See vllm.utils.torch_utils. + fixed = canonicalize_singleton_dim_strides(kv_cache_permute) + if fixed is not kv_cache_permute: + logger.debug( + "Canonicalized degenerate KV cache strides (FlashInfer): " + "shape=%s, strides before=%s, strides after=%s", + kv_cache_permute.shape, + kv_cache_permute.stride(), + fixed.stride(), + ) + kv_cache_permute = fixed # For NVFP4, the kv_cache last dim is full_dim (data + scale packed). # Split into correctly-strided data and scale views. @@ -1568,10 +1584,11 @@ def forward( else: assert isinstance(attn_metadata.prefill, TRTLLMPrefill) # prefill_query may be non-contiguous or have degenerate strides - # First ensure memory contiguity, then fix degenerate strides - # with reshape. contiguous() alone doesn't fix degenerate - # strides when a dimension has size 1. - prefill_query = prefill_query.contiguous().reshape(prefill_query.shape) + # on size=1 dims. contiguous() ensures memory layout; then + # canonicalize_singleton_dim_strides fixes any remaining + # degenerate strides on size=1 dims for TMA alignment. + prefill_query = prefill_query.contiguous() + prefill_query = canonicalize_singleton_dim_strides(prefill_query) workspace_buffer = _get_trtllm_gen_workspace_buffer() block_tables_prefill = attn_metadata.prefill.block_tables seq_lens_prefill = attn_metadata.prefill.seq_lens @@ -1621,11 +1638,9 @@ def forward( # with fp8 kv cache, we can construct a mock block # and mock kv cache with BF16 KV involved in the prefill # - # The inner (block_size, head_size) dims must be - # contiguous; outer dims may have non-canonical strides - # (e.g. cross-layer unified allocation). - # Degenerate strides on outer dims break TMA descriptors - # (see flashinfer-ai/flashinfer#2232). + kv_cache_permute = canonicalize_singleton_dim_strides( + kv_cache_permute + ) kv_strides = kv_cache_permute.stride() assert ( kv_strides[-1] == 1 @@ -1732,12 +1747,13 @@ def forward( if needs_fp8_out: output[:num_decode_tokens].copy_(out_decode.to(output.dtype)) else: - # decode_query may be non-contiguous or have degenerate strides assert isinstance(attn_metadata.decode, TRTLLMDecode) - # First ensure memory contiguity, then fix degenerate strides - # with reshape. contiguous() alone doesn't fix degenerate - # strides when a dimension has size 1. - decode_query = decode_query.contiguous().reshape(decode_query.shape) + # decode_query may be non-contiguous or have degenerate strides + # on size=1 dims. contiguous() ensures memory layout; then + # canonicalize_singleton_dim_strides fixes any remaining + # degenerate strides on size=1 dims for TMA alignment. + decode_query = decode_query.contiguous() + decode_query = canonicalize_singleton_dim_strides(decode_query) workspace_buffer = _get_trtllm_gen_workspace_buffer() block_tables_decode = attn_metadata.decode.block_tables seq_lens_decode = attn_metadata.decode.seq_lens @@ -1748,11 +1764,7 @@ def forward( assert is_strictly_contiguous(workspace_buffer) assert is_strictly_contiguous(block_tables_decode) assert is_strictly_contiguous(seq_lens_decode) - # kv_cache outer dims may be non-contiguous (e.g. - # cross-layer unified allocation), but inner dims - # (block_size, head_size) must be contiguous and - # strides must be canonical to avoid TMA descriptor - # failures (see flashinfer-ai/flashinfer#2232). + kv_cache_permute = canonicalize_singleton_dim_strides(kv_cache_permute) kv_strides = kv_cache_permute.stride() assert ( kv_strides[-1] == 1 and kv_strides[-2] == kv_cache_permute.shape[-1] diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 7b0f00d14c8a..1f0b9bbb19d5 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -403,6 +403,11 @@ def __init__( range(dp_size), local_dp_ranks, placement_groups ): dp_vllm_config = copy.deepcopy(vllm_config) + if dp_size > 1: + # Append the DP rank to instance_id so that per-engine + # identifiers (e.g. Ray actor names in RayExecutorV2) are + # unique across DP replicas. + dp_vllm_config.instance_id = f"{dp_vllm_config.instance_id}_dp{index}" dp_vllm_config.parallel_config.placement_group = pg local_client = index < local_engine_count diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 62912491492e..a02dd62026ad 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -76,6 +76,8 @@ def gumbel_block_argmax( pos_ptr, processed_logits_ptr, processed_logits_stride, + processed_logits_col_ptr, + vocab_size, APPLY_TEMPERATURE: tl.constexpr, ): req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) @@ -88,8 +90,15 @@ def gumbel_block_argmax( if processed_logits_ptr is not None: # Store the temperature-applied logits. + if processed_logits_col_ptr is not None: + col = tl.load(processed_logits_col_ptr) + else: + col = 0 tl.store( - processed_logits_ptr + req_state_idx * processed_logits_stride + block, + processed_logits_ptr + + req_state_idx * processed_logits_stride + + col * vocab_size + + block, logits, mask=mask, ) @@ -121,6 +130,7 @@ def _gumbel_sample_kernel( local_max_stride, processed_logits_ptr, processed_logits_stride, + processed_logits_col_ptr, logits_ptr, logits_stride, expanded_idx_mapping_ptr, @@ -153,6 +163,8 @@ def _gumbel_sample_kernel( pos_ptr, processed_logits_ptr, processed_logits_stride, + processed_logits_col_ptr, + vocab_size, APPLY_TEMPERATURE=APPLY_TEMPERATURE, ) token_id = block_idx * BLOCK_SIZE + idx @@ -167,7 +179,8 @@ def gumbel_sample( seed: torch.Tensor, # [max_num_reqs] pos: torch.Tensor, # [num_tokens] apply_temperature: bool, - processed_logits_out: torch.Tensor | None = None, # [num_reqs, vocab_size] + output_processed_logits: torch.Tensor | None = None, + output_processed_logits_col: torch.Tensor | None = None, ) -> torch.Tensor: num_tokens, vocab_size = logits.shape BLOCK_SIZE = 1024 @@ -179,8 +192,9 @@ def gumbel_sample( local_argmax.stride(0), local_max, local_max.stride(0), - processed_logits_out, - processed_logits_out.stride(0) if processed_logits_out is not None else 0, + output_processed_logits, + output_processed_logits.stride(0) if output_processed_logits is not None else 0, + output_processed_logits_col, logits, logits.stride(0), expanded_idx_mapping, diff --git a/vllm/v1/worker/gpu/sample/prompt_logprob.py b/vllm/v1/worker/gpu/sample/prompt_logprob.py index 11dbf6985279..baa48ebf900c 100644 --- a/vllm/v1/worker/gpu/sample/prompt_logprob.py +++ b/vllm/v1/worker/gpu/sample/prompt_logprob.py @@ -55,10 +55,8 @@ def compute_prompt_logprobs( num_prompt_logprobs = self.num_prompt_logprobs[idx_mapping_np] prompt_lens = prompt_lens[idx_mapping_np] - # NOTE(woosuk): -1 because the last prompt token's hidden state is not - # needed for prompt logprobs. computed_prefill = num_computed_prefill_tokens[idx_mapping_np] - includes_prompt = computed_prefill < prompt_lens - 1 + includes_prompt = computed_prefill < prompt_lens # NOTE(woosuk): If the request was resumed after preemption, its prompt # logprobs must have been computed before preemption. Skip. resumed_after_prompt = prompt_lens < prefill_lens[idx_mapping_np] diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index c6b0aa364f53..efe510f16e22 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -89,9 +89,13 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): dtype=torch.int64, device=device, ) + self.current_draft_step = torch.tensor(0, dtype=torch.int64, device=device) self.last_token_indices = torch.zeros( self.max_num_reqs, dtype=torch.int64, device=device ) + self.arange = torch.arange( + self.max_num_reqs + 1, dtype=torch.int32, device="cpu" + ) self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs( self.draft_model_config @@ -228,9 +232,10 @@ def _sample_draft( logits: torch.Tensor, idx_mapping: torch.Tensor, pos: torch.Tensor, - step: int, + draft_step: torch.Tensor, + draft_logits: torch.Tensor | None, ) -> torch.Tensor: - if self.draft_logits is not None: + if draft_logits is not None: # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise # used for draft and target sampling. return gumbel_sample( @@ -240,7 +245,8 @@ def _sample_draft( self.seeds, pos + 1, apply_temperature=True, - processed_logits_out=self.draft_logits[:, step], + output_processed_logits=draft_logits, + output_processed_logits_col=draft_step, ) else: return logits.argmax(dim=-1) @@ -274,11 +280,63 @@ def prefill( logits, idx_mapping, pos, - step=0, + self.current_draft_step, + self.draft_logits, ) self.hidden_states[:num_reqs] = hidden_states[last_token_indices] self.input_buffers.positions[:num_reqs] = pos + def multi_step_decode( + self, + num_reqs: int, + skip_attn: bool, + batch_desc: BatchExecutionDescriptor, + num_tokens_across_dp: torch.Tensor | None, + ) -> None: + positions = self.input_buffers.positions[:num_reqs] + query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1] + idx_mapping = self.idx_mapping[:num_reqs] + + for step in range(1, self.num_speculative_steps): + attn_metadata = None + slot_mappings_by_layer = None + if not skip_attn: + # Build attention metadata and slot mappings for each draft + # decode step. It is necessary to rebuild the attention + # metadata even when replaying the FULL graph so that any + # attention metadata builder state is updated. + slot_mappings = self.block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + batch_desc.num_tokens, + ) + slot_mappings_by_layer = build_slot_mappings_by_layer( + slot_mappings, self.kv_cache_config + ) + attn_metadata = self._build_draft_attn_metadata( + num_reqs=num_reqs, + num_reqs_padded=batch_desc.num_reqs or num_reqs, + num_tokens_padded=batch_desc.num_tokens, + ) + + # Update the current draft step. + self.current_draft_step.fill_(step) + + # Generate draft tokens for the current step. + if batch_desc.cg_mode == CUDAGraphMode.FULL: + assert self.decode_cudagraph_manager is not None + self.decode_cudagraph_manager.run_fullgraph(batch_desc) + else: + self.generate_draft( + num_reqs, + batch_desc.num_tokens, + attn_metadata, + slot_mappings_by_layer, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=batch_desc.cg_mode, + ) + def generate_draft( self, num_reqs: int, @@ -288,59 +346,52 @@ def generate_draft( num_tokens_across_dp: torch.Tensor | None, cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, ) -> None: - pos = self.input_buffers.positions[:num_reqs] - query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1] idx_mapping = self.idx_mapping[:num_reqs] - for step in range(1, self.num_speculative_steps): - # Run the eagle model. - last_hidden_states, hidden_states = self.run_model( - num_tokens_padded, - attn_metadata, - slot_mappings, - num_tokens_across_dp, - cudagraph_runtime_mode, - ) - last_hidden_states = last_hidden_states[:num_reqs] - hidden_states = hidden_states[:num_reqs] - logits = self.model.compute_logits(last_hidden_states) + positions = self.input_buffers.positions[:num_reqs] + # Run the eagle model forward pass. + last_hidden_states, hidden_states = self.run_model( + num_tokens_padded, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cudagraph_runtime_mode, + ) + last_hidden_states = last_hidden_states[:num_reqs] - draft_tokens = self._sample_draft( - logits, - idx_mapping, - pos, - step=step, - ) - self.draft_tokens[:num_reqs, step] = draft_tokens - - if step < self.num_speculative_steps - 1: - # Update the inputs for the next step. - update_eagle_inputs( - draft_tokens, - hidden_states, - self.input_buffers, - self.hidden_states, - self.max_model_len, - ) - if attn_metadata is not None: - self.block_tables.compute_slot_mappings( - idx_mapping, query_start_loc, pos, num_tokens_padded - ) + # Sample the draft tokens. + logits = self.model.compute_logits(last_hidden_states) + draft_tokens = self._sample_draft( + logits, + idx_mapping, + positions, + self.current_draft_step, + self.draft_logits, + ) + + # Update the inputs for the next step. + update_eagle_draft_inputs( + draft_tokens, + self.current_draft_step, + hidden_states, + self.draft_tokens, + self.hidden_states, + self.input_buffers, + num_reqs, + self.max_model_len, + self.num_speculative_steps, + ) def _build_draft_attn_metadata( self, num_reqs: int, num_reqs_padded: int, num_tokens_padded: int, - max_query_len: int, ) -> dict[str, Any] | None: if not self.draft_attn_layer_names: return None - query_start_loc_cpu = ( - torch.arange(num_reqs_padded + 1, dtype=torch.int32, device="cpu").clamp_( - max=num_reqs - ) - * max_query_len + query_start_loc_cpu = torch.clamp( + self.arange[: num_reqs_padded + 1], max=num_reqs ) block_tables = [ x[:num_reqs_padded] for x in self.block_tables.input_block_tables @@ -354,7 +405,7 @@ def _build_draft_attn_metadata( : num_reqs_padded + 1 ], query_start_loc_cpu=query_start_loc_cpu, - max_query_len=max_query_len, + max_query_len=1, seq_lens=self.input_buffers.seq_lens[:num_reqs_padded], max_seq_len=self.max_model_len, block_tables=block_tables, @@ -373,7 +424,7 @@ def capture( self.last_token_indices.zero_() # Capture the prefill routine (model forward + compute_logits + - # gumbel_sample). + # sample). # For FULL graphs, the entire routine is recorded as one graph. # For PIECEWISE, only the model's compiled regions are captured # and the rest (compute_logits, gumbel_sample) runs eagerly. @@ -387,10 +438,9 @@ def capture( if self.num_speculative_steps == 1: return - # Capture the decode draft generation loop (model forward + - # compute_logits + gumbel_sample + update_eagle_inputs, for - # each step). For FULL graphs, the entire multi-step loop is - # recorded as one graph. + # Capture the decode draft generation routine (model forward + + # compute_logits + sample + update_eagle_inputs) for a single + # step. assert self.decode_cudagraph_manager is not None self.decode_cudagraph_manager.capture( self.generate_draft, @@ -461,9 +511,10 @@ def propose( # Get the input ids and last token indices for the speculator. prepare_eagle_inputs( + self.last_token_indices, + self.current_draft_step, self.input_buffers, input_batch, - self.last_token_indices, num_sampled, num_rejected, last_sampled, @@ -473,12 +524,18 @@ def propose( # When all requests are decoding (no true prefills), each has # num_speculative_steps + 1 tokens, enabling FULL graph replay. - # Mixed or prefill-only batches fall back to PIECEWISE. + uniform_token_count = get_uniform_token_count( + num_reqs, + # Use the actual number of tokens without padding added by + # the target model during FULL cudagraph. + input_batch.num_tokens, + max_query_len, + ) prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( self.prefill_cudagraph_manager, num_reqs, num_tokens, - get_uniform_token_count(num_reqs, num_tokens, max_query_len), + uniform_token_count, dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, @@ -528,48 +585,21 @@ def propose( need_eager=is_profile, ) - attn_metadata_updated = None - slot_mappings_updated = None - if not (dummy_run and skip_attn_for_dummy_run): - # Build attention metadata and slot mappings for the draft - # decode steps. It is necessary to rebuild the attention - # metadata even when replaying the FULL graph so that any - # attention metadata builder state is updated. - slot_mappings = self.block_tables.compute_slot_mappings( - self.idx_mapping[:num_reqs], - self.input_buffers.query_start_loc[: num_reqs + 1], - self.input_buffers.positions[:num_reqs], - decode_batch_desc.num_tokens, - ) - slot_mappings_updated = build_slot_mappings_by_layer( - slot_mappings, self.kv_cache_config - ) - attn_metadata_updated = self._build_draft_attn_metadata( - num_reqs=num_reqs, - num_reqs_padded=decode_batch_desc.num_reqs or num_reqs, - num_tokens_padded=decode_batch_desc.num_tokens, - max_query_len=1, - ) + # Generate the remaining num_speculative_steps - 1 draft tokens. + self.multi_step_decode( + num_reqs, + dummy_run and skip_attn_for_dummy_run, + decode_batch_desc, + num_tokens_across_dp, + ) - if decode_batch_desc.cg_mode == CUDAGraphMode.FULL: - # Replay the full graph for draft generation. - assert self.decode_cudagraph_manager is not None - self.decode_cudagraph_manager.run_fullgraph(decode_batch_desc) - else: - self.generate_draft( - num_reqs, - decode_batch_desc.num_tokens, - attn_metadata_updated, - slot_mappings_updated, - num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=decode_batch_desc.cg_mode, - ) return self.draft_tokens[:num_reqs] @triton.jit def _prepare_eagle_inputs_kernel( last_token_indices_ptr, + eagle_current_draft_step_ptr, eagle_input_ids_ptr, eagle_positions_ptr, eagle_query_start_loc_ptr, @@ -630,6 +660,8 @@ def _prepare_eagle_inputs_kernel( # Copy sequence lengths. tl.store(eagle_seq_lens_ptr + req_idx, seq_len) if req_idx == (num_reqs - 1): + # Reset the current draft step to 0. + tl.store(eagle_current_draft_step_ptr, 0) # Pad query_start_loc for CUDA graphs. for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE): block = i + tl.arange(0, BLOCK_SIZE) @@ -648,10 +680,11 @@ def _prepare_eagle_inputs_kernel( def prepare_eagle_inputs( - input_buffers: InputBuffers, - input_batch: InputBatch, # [num_reqs] last_token_indices: torch.Tensor, + current_draft_step: torch.Tensor, + input_buffers: InputBuffers, + input_batch: InputBatch, # [num_reqs] num_sampled: torch.Tensor, # [num_reqs] @@ -665,6 +698,7 @@ def prepare_eagle_inputs( num_reqs = input_batch.num_reqs _prepare_eagle_inputs_kernel[(num_reqs,)]( last_token_indices, + current_draft_step, input_buffers.input_ids, input_buffers.positions, input_buffers.query_start_loc, @@ -685,7 +719,7 @@ def prepare_eagle_inputs( @triton.jit -def _prepare_eagle_docode_kernel( +def _prepare_eagle_decode_kernel( draft_tokens_ptr, draft_tokens_stride, target_seq_lens_ptr, @@ -742,7 +776,7 @@ def prepare_eagle_decode( max_num_reqs: int, ): num_reqs = draft_tokens.shape[0] - _prepare_eagle_docode_kernel[(num_reqs + 1,)]( + _prepare_eagle_decode_kernel[(num_reqs + 1,)]( draft_tokens, draft_tokens.stride(0), target_seq_lens, @@ -758,36 +792,55 @@ def prepare_eagle_decode( @triton.jit -def _update_eagle_inputs_kernel( +def _update_eagle_draft_inputs_kernel( + output_draft_tokens_ptr, + output_draft_tokens_stride, + next_input_hidden_states_ptr, + next_input_hidden_states_stride, input_ids_ptr, positions_ptr, - input_hidden_states_ptr, - input_hidden_states_stride, seq_lens_ptr, - max_model_len, draft_tokens_ptr, - output_hidden_states_ptr, - output_hidden_states_stride, + current_draft_step_ptr, + hidden_states_ptr, + hidden_states_stride, hidden_size, + max_model_len, + num_speculative_steps, BLOCK_SIZE: tl.constexpr, ): req_idx = tl.program_id(0) - # Draft token -> Input ID. + # Write the sampled draft token into self.draft_tokens[req_idx, step]. draft_token = tl.load(draft_tokens_ptr + req_idx) + step = tl.load(current_draft_step_ptr) + tl.store( + output_draft_tokens_ptr + req_idx * output_draft_tokens_stride + step, + draft_token, + ) + + if step >= num_speculative_steps - 1: + # This is the final step. Skip updating draft forward inputs. + return + + # Write the sampled draft token into the input ids tensor for the next + # forward pass. tl.store(input_ids_ptr + req_idx, draft_token) - # Output hidden states -> Input hidden states. + # Copy hidden states into the input hidden states tensor for the next + # forward pass. for i in range(0, hidden_size, BLOCK_SIZE): block = i + tl.arange(0, BLOCK_SIZE) mask = block < hidden_size - output_hidden_states = tl.load( - output_hidden_states_ptr + req_idx * output_hidden_states_stride + block, + hidden_states = tl.load( + hidden_states_ptr + req_idx * hidden_states_stride + block, mask=mask, ) tl.store( - input_hidden_states_ptr + req_idx * input_hidden_states_stride + block, - output_hidden_states, + next_input_hidden_states_ptr + + req_idx * next_input_hidden_states_stride + + block, + hidden_states, mask=mask, ) @@ -803,24 +856,32 @@ def _update_eagle_inputs_kernel( tl.store(seq_lens_ptr + req_idx, seq_len) -def update_eagle_inputs( +def update_eagle_draft_inputs( draft_tokens: torch.Tensor, - output_hidden_states: torch.Tensor, - input_buffers: InputBuffers, + current_draft_step: torch.Tensor, hidden_states: torch.Tensor, + output_draft_tokens: torch.Tensor, + next_input_hidden_states: torch.Tensor, + input_buffers: InputBuffers, + num_reqs: int, max_model_len: int, + num_speculative_steps: int, ): - num_reqs, hidden_size = output_hidden_states.shape - _update_eagle_inputs_kernel[(num_reqs,)]( + _, hidden_size = hidden_states.shape + _update_eagle_draft_inputs_kernel[(num_reqs,)]( + output_draft_tokens, + output_draft_tokens.stride(0), + next_input_hidden_states, + next_input_hidden_states.stride(0), input_buffers.input_ids, input_buffers.positions, - hidden_states, - hidden_states.stride(0), input_buffers.seq_lens, - max_model_len, draft_tokens, - output_hidden_states, - output_hidden_states.stride(0), + current_draft_step, + hidden_states, + hidden_states.stride(0), hidden_size, + max_model_len, + num_speculative_steps, BLOCK_SIZE=1024, ) diff --git a/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py index 9d86372e624b..10b29433efb2 100644 --- a/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py @@ -392,8 +392,10 @@ def _resample_kernel( temp_ptr, seed_ptr, pos_ptr, - None, - 0, + None, # processed_logits_ptr + 0, # processed_logits_stride + None, # processed_logits_col_ptr + vocab_size, APPLY_TEMPERATURE=False, ) token_id = block_idx * BLOCK_SIZE + idx diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 75898f463272..44e0efaaa2f2 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -49,6 +49,8 @@ class CachedRequestState: lora_request: LoRARequest | None = None prompt_embeds: torch.Tensor | None = None + # To accumulate prompt logprobs tensor chunks across prefill steps. + in_progress_prompt_logprobs_cpu: LogprobsTensors | None = None # Per-position mask for mixed-mode inputs (e.g chat completion with # prompt_embeds content parts). See `Request.prompt_is_token_ids`. @@ -255,9 +257,6 @@ def __init__( # More efficient than num_logprobs=-1 when only a few tokens are needed self.logprob_token_ids: dict[str, list[int]] = {} - # To accumulate prompt logprobs tensor chunks across prefill steps. - self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {} - # Internal representation of per-step batch state changes, used for # reordering persistent batch and generating logitsprocs batch state # updates. Should reset each step. @@ -552,7 +551,6 @@ def remove_request(self, req_id: str) -> int | None: self.generators.pop(req_index, None) self.num_logprobs.pop(req_id, None) self.logprob_token_ids.pop(req_id, None) - self.in_progress_prompt_logprobs_cpu.pop(req_id, None) if self.prev_req_id_to_index is not None: self.prev_req_id_to_index.pop(req_id, None) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index bcab2ca2d4c2..0ca530c15bac 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5094,7 +5094,6 @@ def _get_prompt_logprobs_dict( if not num_prompt_logprobs_dict: return {} - in_progress_dict = self.input_batch.in_progress_prompt_logprobs_cpu prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {} # Since prompt logprobs are a rare feature, prioritize simple, @@ -5118,14 +5117,14 @@ def _get_prompt_logprobs_dict( ) # Set up target LogprobsTensors object. - logprobs_tensors = in_progress_dict.get(req_id) - if not logprobs_tensors: + logprobs_tensors = request.in_progress_prompt_logprobs_cpu + if logprobs_tensors is None: # Create empty logprobs CPU tensors for the entire prompt. # If chunked, we'll copy in slice by slice. logprobs_tensors = LogprobsTensors.empty_cpu( num_prompt_tokens - 1, num_prompt_logprobs + 1 ) - in_progress_dict[req_id] = logprobs_tensors + request.in_progress_prompt_logprobs_cpu = logprobs_tensors # Determine number of logits to retrieve. start_idx = request.num_computed_tokens @@ -5182,7 +5181,7 @@ def _get_prompt_logprobs_dict( # num_prompt_logprobs_dict. for req_id in completed_prefill_reqs: del num_prompt_logprobs_dict[req_id] - del in_progress_dict[req_id] + self.requests[req_id].in_progress_prompt_logprobs_cpu = None # Must synchronize the non-blocking GPU->CPU transfers. if prompt_logprobs_dict: