karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

Merge branch 'main' into wentao-fix-mypy-models-a-b

Browse Source
This commit is contained in:
Wentao Ye
2026-07-23 14:29:26 -04:00
committed by GitHub
parent 102a9b719c b354734d17
commit e052f3a44f
28 changed files with 352 additions and 224 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh
+7 -6
View File
@@ -1,10 +1,11 @@
#!/bin/bash
set -euox pipefail
export VLLM_CPU_KVCACHE_SPACE=1
export VLLM_CPU_KVCACHE_SPACE=1
export VLLM_CPU_CI_ENV=1
# Reduce sub-processes for acceleration
export TORCH_COMPILE_DISABLE=1
# Skip torch.compile via vLLM's --enforce-eager flag (passed below) instead of
# TORCH_COMPILE_DISABLE=1, which torch 2.12 no longer treats as a silent no-op
# when callers specify fullgraph=True.
export VLLM_ENABLE_V1_MULTIPROCESSING=0
SDE_ARCHIVE="sde-external-10.7.0-2026-02-18-lin.tar.xz"
@@ -49,15 +50,15 @@ wait_for_pid_and_check_log() {
}
# Test Sky Lake (AVX512F)
./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_0.log 2>&1 &
./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_0.log 2>&1 &
PID_TEST_0=$!
# Test Cascade Lake (AVX512F + VNNI)
./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_1.log 2>&1 &
./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_1.log 2>&1 &
PID_TEST_1=$!
# Test Cooper Lake (AVX512F + VNNI + BF16)
./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_2.log 2>&1 &
./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_2.log 2>&1 &
PID_TEST_2=$!
wait_for_pid_and_check_log $PID_TEST_0 test_0.log
.buildkite/test_areas/pytorch.yaml
+1 -1
View File
@@ -5,7 +5,7 @@ steps:
- label: PyTorch Compilation Unit Tests
device: h200_35gb
key: pytorch-compilation-unit-tests
timeout_in_minutes: 110
timeout_in_minutes: 150
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
CMakeLists.txt
+2 -2
View File
@@ -68,8 +68,8 @@ endif()
# requirements.txt files and should be kept consistent. The ROCm torch
# versions are derived from docker/Dockerfile.rocm
#
set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0")
set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0")
set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0")
set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0")
# TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported-
# version check would always warn. Only treat it as a nightly build when the
# value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which
docker/Dockerfile
+1 -1
View File
@@ -22,7 +22,7 @@
# docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
# =============================================================================
ARG CUDA_VERSION=13.0.2
ARG CUDA_VERSION=13.0.3
ARG PYTHON_VERSION=3.12
ARG UBUNTU_VERSION=22.04
docker/versions.json
+3 -3
View File
@@ -2,7 +2,7 @@
"_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
"variable": {
"CUDA_VERSION": {
"default": "13.0.2"
"default": "13.0.3"
},
"PYTHON_VERSION": {
"default": "3.12"
@@ -11,10 +11,10 @@
"default": "22.04"
},
"BUILD_BASE_IMAGE": {
"default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
"default": "nvidia/cuda:13.0.3-devel-ubuntu22.04"
},
"FINAL_BASE_IMAGE": {
"default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
"default": "nvidia/cuda:13.0.3-base-ubuntu22.04"
},
"BUILD_OS": {
"default": "ubuntu"
pyproject.toml
+1 -1
View File
@@ -7,7 +7,7 @@ requires = [
"setuptools>=77.0.3,<81.0.0",
"setuptools-scm>=8.0",
"setuptools-rust>=1.9.0",
"torch == 2.11.0",
"torch == 2.13.0",
"wheel",
"jinja2",
]
requirements/build/cpu.txt
+2 -2
View File
@@ -4,8 +4,8 @@ packaging>=24.2
setuptools==77.0.3 # this version can reuse CMake build dir
setuptools-scm>=8
setuptools-rust>=1.9.0
torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
torch==2.11.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
torch==2.13.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
torch==2.13.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
wheel
jinja2>=3.1.6
regex
requirements/build/cuda.txt
+1 -1
View File
@@ -5,7 +5,7 @@ packaging>=24.2
setuptools>=77.0.3,<81.0.0
setuptools-scm>=8
setuptools-rust>=1.9.0
torch==2.11.0
torch==2.13.0
wheel
jinja2>=3.1.6
regex
requirements/cpu.txt
+2 -2
View File
@@ -6,8 +6,8 @@ setuptools==77.0.3 # this version can reuse CMake build dir
numba == 0.65.0; platform_machine != "s390x" # Required for N-gram speculative decoding
# Dependencies for CPUs
torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
torch==2.11.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
torch==2.13.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
torch==2.13.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
# required for the image processor of minicpm-o-2_6, this must be updated alongside torch
torchaudio; platform_machine != "s390x" and platform_machine != "riscv64"
requirements/cuda.txt
+2 -2
View File
@@ -4,10 +4,10 @@
numba == 0.65.0 # Required for N-gram speculative decoding
# Dependencies for NVIDIA GPUs
torch==2.11.0
torch==2.13.0
torchaudio==2.11.0
# These must be updated alongside torch
torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
torchvision==0.28.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
torchcodec >= 0.14
PyNvVideoCodec==2.0.4
# FlashInfer should be updated together with the Dockerfile
requirements/test/cpu.txt
+2 -2
View File
@@ -1107,7 +1107,7 @@ tokenizers==0.22.2
# -r requirements/test/../common.txt
# -r requirements/test/cuda.in
# transformers
torch==2.11.0+cpu
torch==2.13.0+cpu
# via
# -r requirements/test/cuda.in
# accelerate
@@ -1134,7 +1134,7 @@ torchaudio==2.11.0+cpu
# vocos
torchcodec==0.14.0+cpu
# via -r requirements/test/cuda.in
torchvision==0.26.0+cpu
torchvision==0.28.0+cpu
# via
# -r requirements/test/cuda.in
# open-clip-torch
requirements/test/cuda.in
+2 -2
View File
@@ -28,9 +28,9 @@ soundfile # required for audio tests
jiwer # required for audio tests
tblib # for pickling test exceptions
timm >=1.0.17 # required for internvl and gemma3n-mm test
torch==2.11.0
torch==2.13.0
torchaudio==2.11.0
torchvision==0.26.0
torchvision==0.28.0
transformers_stream_generator # required for qwen-vl test
matplotlib # required for qwen-vl test
mistral_common[image,audio] >= 1.11.5 # required for voxtral test
requirements/test/cuda.txt
+11 -9
View File
@@ -159,7 +159,7 @@ cuda-bindings==13.0.3
# via torch
cuda-pathfinder==1.3.3
# via cuda-bindings
cuda-toolkit==13.0.2
cuda-toolkit==13.0.3.0
# via torch
cupy-cuda12x==13.6.0
# via ray
@@ -599,7 +599,7 @@ numpy==2.2.6
# tritonclient
# vocos
# xgrammar
nvidia-cublas==13.1.0.3
nvidia-cublas==13.1.1.3
# via
# cuda-toolkit
# nvidia-cudnn-cu13
@@ -607,10 +607,12 @@ nvidia-cublas==13.1.0.3
nvidia-cuda-cupti==13.0.85
# via cuda-toolkit
nvidia-cuda-nvrtc==13.0.88
# via cuda-toolkit
# via
# cuda-toolkit
# nvidia-cublas
nvidia-cuda-runtime==13.0.96
# via cuda-toolkit
nvidia-cudnn-cu13==9.19.0.56
nvidia-cudnn-cu13==9.20.0.48
# via torch
nvidia-cufft==12.0.0.61
# via cuda-toolkit
@@ -624,9 +626,9 @@ nvidia-cusparse==12.6.3.3
# via
# cuda-toolkit
# nvidia-cusolver
nvidia-cusparselt-cu13==0.8.0
nvidia-cusparselt-cu13==0.8.1
# via torch
nvidia-nccl-cu13==2.28.9
nvidia-nccl-cu13==2.29.7
# via torch
nvidia-nvjitlink==13.0.88
# via
@@ -1202,7 +1204,7 @@ tokenizers==0.22.2
# -r requirements/test/../common.txt
# -r requirements/test/cuda.in
# transformers
torch==2.11.0+cu130
torch==2.13.0+cu130
# via
# -c requirements/cuda.txt
# -r requirements/test/cuda.in
@@ -1233,7 +1235,7 @@ torchcodec==0.14.0+cu130
# via
# -c requirements/cuda.txt
# -r requirements/test/cuda.in
torchvision==0.26.0+cu130
torchvision==0.28.0+cu130
# via
# -c requirements/cuda.txt
# -r requirements/test/cuda.in
@@ -1270,7 +1272,7 @@ transformers==5.13.1
# xgrammar
transformers-stream-generator==0.0.5
# via -r requirements/test/cuda.in
triton==3.6.0
triton==3.7.1
# via
# torch
# xgrammar
tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py
+13 -2
View File
@@ -142,8 +142,19 @@ def qwen2audio_aligned_content_and_embeds_b64() -> tuple[str, str]:
@pytest.mark.asyncio
@pytest.mark.parametrize(
"audio_first",
[True, False],
ids=["audio_embeds-then-text", "text-then-audio_embeds"],
[
pytest.param(True, id="audio_embeds-then-text"),
pytest.param(
False,
id="text-then-audio_embeds",
marks=pytest.mark.xfail(
reason="torch 2.12 regression: prompt_embeds output diverges "
"from raw-text when text precedes audio; "
"https://github.com/pytorch/pytorch/issues/184431",
strict=True,
),
),
],
)
async def test_text_content_and_prompt_embeds_match_with_audio_embeds(
qwen2audio_client: openai.AsyncOpenAI,
tests/kernels/moe/test_block_int8.py
+2 -2
View File
@@ -82,9 +82,9 @@ def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
).sum(dim=1)
@pytest.fixture(autouse=True, scope="module")
@pytest.fixture(autouse=True)
def setup_cuda():
"""Sets the default CUDA device for all tests in this module."""
"""Sets the default CUDA device before each test in this module."""
torch.set_default_device("cuda")
tests/kernels/moe/test_triton_moe_ptpc_fp8.py
+2 -2
View File
@@ -102,9 +102,9 @@ def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk):
).sum(dim=1)
@pytest.fixture(autouse=True, scope="module")
@pytest.fixture(autouse=True)
def setup_cuda():
"""Sets the default CUDA device for all tests in this module."""
"""Sets the default CUDA device before each test in this module."""
torch.set_default_device("cuda")
tests/kernels/quantization/test_block_int8.py
+10 -10
View File
@@ -28,12 +28,6 @@ BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
@pytest.fixture(autouse=True, scope="module")
def setup_cuda():
"""Sets the default CUDA device for all tests in this module."""
torch.set_default_device("cuda")
@pytest.mark.parametrize(
"M,N,K,block_size,out_dtype,seed",
itertools.product(M, N, K, BLOCK_SIZE, DTYPES, SEEDS),
@@ -41,22 +35,28 @@ def setup_cuda():
@torch.inference_mode()
def test_w8a8_block_int8_matmul(M, N, K, block_size, out_dtype, seed):
torch.manual_seed(seed)
device = current_platform.device_type
factor_for_scale = 1e-2
int8_info = torch.iinfo(torch.int8)
int8_max, int8_min = int8_info.max, int8_info.min
A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * int8_max
A_fp32 = torch.rand(M, K, dtype=torch.float32, device=device)
A_fp32 = (A_fp32 - 0.5) * 2 * int8_max
A_fp8 = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn)
B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * int8_max
B_fp32 = torch.rand(N, K, dtype=torch.float32, device=device)
B_fp32 = (B_fp32 - 0.5) * 2 * int8_max
B_fp8 = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn)
block_n, block_k = block_size[0], block_size[1]
n_tiles = (N + block_n - 1) // block_n
k_tiles = (K + block_k - 1) // block_k
As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale
Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
As = torch.rand(M, k_tiles, dtype=torch.float32, device=device) * factor_for_scale
Bs = (
torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device)
* factor_for_scale
)
ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
out = w8a8_block_int8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
tests/kernels/quantization/test_int8_kernel.py
+12 -10
View File
@@ -82,12 +82,6 @@ def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight, topk_ids
).sum(dim=1)
@pytest.fixture(autouse=True, scope="module")
def setup_cuda():
"""Sets the default CUDA device for all tests in this module."""
torch.set_default_device("cuda")
DTYPES = [torch.half, torch.bfloat16]
M = [1, 33]
N = [128, 1024]
@@ -104,6 +98,7 @@ SEEDS = [0]
@torch.inference_mode()
def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed):
torch.manual_seed(seed)
device = current_platform.device_type
# Initialize int8 quantization parameters
factor_for_scale = 1e-2
int8_max = 127
@@ -111,19 +106,26 @@ def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed):
# Input tensor
# M * K
a = torch.randn((M, K), dtype=dtype) / 10
a = torch.randn((M, K), dtype=dtype, device=device) / 10
# Generate int8 weights
w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2
w1_fp32 = (
torch.rand(
(E, 2 * N, K),
dtype=torch.float32,
device=device,
)
- 0.5
) * 2
w1 = (w1_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2
w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32, device=device) - 0.5) * 2
w2 = (w2_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
# Generate scale for each column (per-column quantization)
w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * factor_for_scale
w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale
score = torch.randn((M, E), dtype=dtype)
score = torch.randn((M, E), dtype=dtype, device=device)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weights, topk_ids = torch.topk(score, topk)
tests/models/multimodal/generation/test_common.py
+9 -1
View File
@@ -908,7 +908,15 @@ VLM_TEST_SETTINGS = {
multi_image_prompt="Picture 1: <vlm_image>\nPicture 2: <vlm_image>\nDescribe these two images with one paragraph respectively.", # noqa: E501
max_model_len=4096,
max_num_seqs=2,
num_logprobs=10,
# torch 2.13 accumulates CPU numerical drift in the qwen2_vl multi-image
# path: HF and vLLM agree for a long prefix (~69 tokens) then a token
# flips outside vLLM's top-N only near the end of the generation. The
# window is already at the max_logprobs=20 cap, so widening it further is
# not possible. Treat this as acceptable drift and cap max_tokens on CPU
# so the compared prefix stays before the divergence, keeping the
# multi-image path under test. See pytorch/pytorch#187735.
max_tokens=64 if current_platform.is_cpu() else 128,
num_logprobs=20 if current_platform.is_cpu() else 10,
auto_cls=AutoModelForImageTextToText,
vllm_output_post_proc=model_utils.qwen2_vllm_to_hf_output,
image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
tests/v1/cudagraph/test_cudagraph_mode.py
+3 -24
View File
@@ -1,11 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref
from contextlib import ExitStack
import pytest
from tests.utils import wait_for_gpu_memory_to_clear
from tests.utils import create_new_process_for_each_test
from tests.v1.attention.utils import full_cg_backend_configs as backend_configs
from vllm import LLM
from vllm.config import CompilationConfig, CompilationMode
@@ -32,6 +31,7 @@ else:
@pytest.mark.parametrize("backend_name, cudagraph_mode, supported", combo_cases_1)
@create_new_process_for_each_test("spawn")
def test_backend_and_cudagraph_mode_combo(backend_name, cudagraph_mode, supported):
if backend_name == "FlashInfer":
try:
@@ -64,17 +64,6 @@ def test_backend_and_cudagraph_mode_combo(backend_name, cudagraph_mode, supporte
),
)
llm.generate(["Hello, my name is"] * 10)
# when above code raises, `llm` may be undefined, so we need to catch that
try:
llm = weakref.proxy(llm)
del llm
except UnboundLocalError:
pass
wait_for_gpu_memory_to_clear(
devices=[0],
threshold_ratio=0.1,
)
# test cudagraph_mode with different compilation mode.
@@ -98,6 +87,7 @@ combo_cases_2 = [
@pytest.mark.parametrize(
"backend_name,cudagraph_mode,compilation_mode,supported", combo_cases_2
)
@create_new_process_for_each_test("spawn")
def test_cudagraph_compilation_combo(
backend_name, cudagraph_mode, compilation_mode, supported
):
@@ -120,14 +110,3 @@ def test_cudagraph_compilation_combo(
),
)
llm.generate(["Hello, my name is"] * 10)
# when above code raises, `llm` may be undefined, so we need to catch that
try:
llm = weakref.proxy(llm)
del llm
except UnboundLocalError:
pass
finally:
wait_for_gpu_memory_to_clear(
devices=[0],
threshold_ratio=0.1,
)
tests/v1/determinism/test_rms_norm_batch_invariant.py
+35 -45
View File
@@ -1,11 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Test batch-invariant RMS normalization against standard implementations.
This test compares the Triton-based batch-invariant RMS norm implementation
with the standard CUDA-based implementation to ensure numerical accuracy.
"""
"""Test batch-invariant RMS normalization against a PyTorch reference."""
import pytest
import torch
@@ -14,18 +9,30 @@ from utils import skip_if_not_cuda, skip_unsupported
from vllm.model_executor.layers.batch_invariant import (
rms_norm_batch_invariant,
)
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.platforms import current_platform
DEVICE_TYPE = current_platform.device_type
def _rms_norm_reference(
input_tensor: torch.Tensor,
weight: torch.Tensor,
eps: float,
) -> torch.Tensor:
"""Compute RMSNorm independently using PyTorch operations."""
input_fp32 = input_tensor.float()
output = input_fp32 * torch.rsqrt(
input_fp32.square().mean(dim=-1, keepdim=True) + eps
)
return (output * weight.float()).to(input_tensor.dtype)
@skip_if_not_cuda
@pytest.mark.parametrize("batch_size", [1, 4, 16, 64])
@pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("eps", [1e-6, 1e-5])
def test_rms_norm_batch_invariant_vs_standard(
def test_rms_norm_batch_invariant_vs_reference(
default_vllm_config,
batch_size: int,
hidden_size: int,
@@ -33,11 +40,10 @@ def test_rms_norm_batch_invariant_vs_standard(
eps: float,
):
"""
Compare batch-invariant Triton RMS norm against standard CUDA implementation.
Compare batch-invariant Triton RMS norm against a PyTorch reference.
Tests that the Triton-based batch-invariant RMS norm produces numerically
equivalent results to the standard CUDA implementation across various
configurations.
equivalent results to an independent implementation across configurations.
"""
device = torch.device(DEVICE_TYPE)
@@ -46,11 +52,7 @@ def test_rms_norm_batch_invariant_vs_standard(
input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
weight = torch.randn(hidden_size, dtype=dtype, device=device)
# Standard implementation (CUDA ops)
rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
rms_norm_layer.weight.data = weight.clone()
standard_output = rms_norm_layer.forward_cuda(input_tensor)
reference_output = _rms_norm_reference(input_tensor, weight, eps)
# Batch-invariant implementation (Triton)
triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)
@@ -64,7 +66,7 @@ def test_rms_norm_batch_invariant_vs_standard(
torch.testing.assert_close(
triton_output,
standard_output,
reference_output,
rtol=rtol,
atol=atol,
msg=f"RMS norm mismatch for batch_size={batch_size}, "
@@ -127,7 +129,7 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(
)
merged_single = x_single + residual_single
ref_out = rms_norm_batch_invariant(merged_single, weight, eps=eps)
ref_out = _rms_norm_reference(merged_single, weight, eps)
torch.testing.assert_close(
residual_out_single,
@@ -162,7 +164,7 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(
rtol=rtol,
atol=atol,
msg="Fused add RMSNorm output should stay numerically close to the "
"batch-invariant RMSNorm reference",
"PyTorch RMSNorm reference",
)
@@ -189,10 +191,7 @@ def test_rms_norm_3d_input(
)
weight = torch.randn(hidden_size, dtype=dtype, device=device)
# Standard implementation
rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
rms_norm_layer.weight.data = weight.clone()
standard_output = rms_norm_layer.forward_cuda(input_tensor)
reference_output = _rms_norm_reference(input_tensor, weight, eps)
# Batch-invariant implementation
triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)
@@ -202,7 +201,7 @@ def test_rms_norm_3d_input(
torch.testing.assert_close(
triton_output,
standard_output,
reference_output,
rtol=rtol,
atol=atol,
msg=f"RMS norm mismatch for 3D input with batch_size={batch_size}, "
@@ -238,20 +237,17 @@ def test_rms_norm_numerical_stability(default_vllm_config):
weight = torch.ones(hidden_size, dtype=dtype, device=device)
for idx, input_tensor in enumerate(test_cases):
# Standard implementation
rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
rms_norm_layer.weight.data = weight.clone()
standard_output = rms_norm_layer.forward_cuda(input_tensor)
reference_output = _rms_norm_reference(input_tensor, weight, eps)
# Batch-invariant implementation
triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)
# Check for NaN or Inf
assert not torch.isnan(standard_output).any(), (
f"Standard RMS norm produced NaN for test case {idx}"
assert not torch.isnan(reference_output).any(), (
f"Reference RMS norm produced NaN for test case {idx}"
)
assert not torch.isinf(standard_output).any(), (
f"Standard RMS norm produced Inf for test case {idx}"
assert not torch.isinf(reference_output).any(), (
f"Reference RMS norm produced Inf for test case {idx}"
)
assert not torch.isnan(triton_output).any(), (
f"Triton RMS norm produced NaN for test case {idx}"
@@ -263,7 +259,7 @@ def test_rms_norm_numerical_stability(default_vllm_config):
# Compare outputs - very lenient for extreme values with float16
torch.testing.assert_close(
triton_output,
standard_output,
reference_output,
rtol=2e-1, # 20% tolerance for extreme values
atol=2e-1,
msg=f"RMS norm mismatch for extreme value test case {idx}",
@@ -321,10 +317,7 @@ def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):
input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
weight = torch.randn(hidden_size, dtype=dtype, device=device)
# Standard implementation
rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
rms_norm_layer.weight.data = weight.clone()
standard_output = rms_norm_layer.forward_cuda(input_tensor)
reference_output = _rms_norm_reference(input_tensor, weight, eps)
# Batch-invariant implementation
triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)
@@ -334,7 +327,7 @@ def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):
torch.testing.assert_close(
triton_output,
standard_output,
reference_output,
rtol=rtol,
atol=atol,
msg=f"RMS norm mismatch for hidden_size={hidden_size}",
@@ -420,21 +413,18 @@ if __name__ == "__main__":
input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
weight = torch.randn(hidden_size, dtype=dtype, device=device)
# Standard implementation
rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
rms_norm_layer.weight.data = weight.clone()
standard_output = rms_norm_layer.forward_cuda(input_tensor)
reference_output = _rms_norm_reference(input_tensor, weight, eps)
# Batch-invariant implementation
triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)
# Compare
max_diff = (triton_output - standard_output).abs().max().item()
mean_diff = (triton_output - standard_output).abs().mean().item()
max_diff = (triton_output - reference_output).abs().max().item()
mean_diff = (triton_output - reference_output).abs().mean().item()
print(f"Max difference: {max_diff:.6e}")
print(f"Mean difference: {mean_diff:.6e}")
print(f"Standard output sample: {standard_output[0, :5].tolist()}")
print(f"Reference output sample: {reference_output[0, :5].tolist()}")
print(f"Triton output sample: {triton_output[0, :5].tolist()}")
if max_diff < 1e-3:
tests/v1/kv_connector/nixl_integration/test_nixl_imports.py
+11 -1
View File
@@ -61,7 +61,17 @@ def test_nixl_and_nixl_ep_imports() -> None:
importlib.import_module("nixl._bindings")
# Exercise the NIXL EP extension used by fused MoE expert parallelism.
nixl_ep = importlib.import_module("nixl_ep")
try:
nixl_ep = importlib.import_module("nixl_ep")
except ImportError as e:
if "materialize_cow_storage" in str(e) or "undefined symbol" in str(e):
pytest.xfail(
"nixl_ep prebuilt extension is ABI-incompatible with this torch "
"(undefined symbol c10::impl::cow::materialize_cow_storage); "
"needs a nixl rebuild against torch 2.13. "
"See pytorch/pytorch#187727 and ai-dynamo/nixl#1798."
)
raise
print(f"nixl_ep: {nixl_ep.__file__}")
assert nixl_ep.__file__ is not None
tests/v1/structured_output/test_reasoning_structured_output.py
+119 -27
View File
@@ -73,6 +73,7 @@ class TestReasoningStructuredOutput:
request.all_token_ids = [1, 2, 3, 4, 5, 6, 7, 8]
request.num_computed_tokens = 5
request.num_output_placeholders = 0
request.request_id = "mock_req"
return request
@pytest.fixture
@@ -208,37 +209,11 @@ class TestReasoningStructuredOutput:
mock_request_with_structured_output
)
# Should set reasoning_ended to True but return False for this step
# The scheduler trims the reasoning prefix before advancing the grammar.
assert (
mock_request_with_structured_output.structured_output_request.reasoning_ended
is True
)
assert result is False
def test_should_advance_reasoning_just_ended_with_spec_decode_structural_tag(
self,
manager_with_reasoner,
mock_request_with_structured_output,
):
"""When reasoning ends this step, advance immediately for structural
tags with speculative decoding."""
structured_req = mock_request_with_structured_output.structured_output_request
structured_req.reasoning_ended = False
structured_req.structured_output_key = (
StructuredOutputOptions.STRUCTURAL_TAG,
"{}",
)
reasoner = MockReasoner(tokenizer=Mock())
reasoner.is_reasoning_end_streaming.return_value = True
structured_req.reasoner = reasoner
manager_with_reasoner.vllm_config.speculative_config = Mock()
result = manager_with_reasoner.should_advance(
mock_request_with_structured_output
)
assert structured_req.reasoning_ended is True
assert result is True
def test_should_advance_reasoning_already_ended(
@@ -258,3 +233,120 @@ class TestReasoningStructuredOutput:
# Should return True since reasoning has ended
assert result is True
def test_should_advance_uses_new_token_ids_when_provided(
self,
manager_with_reasoner,
mock_request_with_structured_output,
):
"""Regression for #43388: when caller passes new_token_ids, the
reasoner sees the exact multi-token delta rather than the
placeholder-derived window.
"""
structured_req = mock_request_with_structured_output.structured_output_request
structured_req.reasoning_ended = False
end_token_id = 248069
reasoner = MockReasoner(tokenizer=Mock())
# Detection mirrors the real Qwen3 parser: end token in the delta.
reasoner.is_reasoning_end_streaming = Mock(
side_effect=lambda input_ids, delta_ids: end_token_id in list(delta_ids)
)
structured_req.reasoner = reasoner
# Scenario from #43388: async + spec decode K=4, 4 tokens accepted
# but only 1 placeholder remains (some drafts were rejected).
# The placeholder math would yield delta=[271] and miss </think>.
# Passing new_token_ids must override that.
new_token_ids = [9, 198, end_token_id, 271]
mock_request_with_structured_output.all_token_ids = [
1,
2,
3,
4,
5,
] + new_token_ids
mock_request_with_structured_output.num_computed_tokens = 9
mock_request_with_structured_output.num_output_placeholders = 1
result = manager_with_reasoner.should_advance(
mock_request_with_structured_output,
new_token_ids=new_token_ids,
)
# First call to is_reasoning_end_streaming was with the full
# new_token_ids (not the truncated placeholder window).
first_call = reasoner.is_reasoning_end_streaming.call_args_list[0]
_, called_delta = first_call.args
assert list(called_delta) == new_token_ids
assert structured_req.reasoning_ended is True
assert result is True
def test_should_advance_without_new_token_ids_falls_back(
self,
manager_with_reasoner,
mock_request_with_structured_output,
):
"""Backward compat: callers that don't pass new_token_ids keep
the original placeholder-derived delta window.
"""
structured_req = mock_request_with_structured_output.structured_output_request
structured_req.reasoning_ended = False
reasoner = MockReasoner(tokenizer=Mock())
reasoner.is_reasoning_end_streaming.return_value = False
structured_req.reasoner = reasoner
mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5]
mock_request_with_structured_output.num_computed_tokens = 5
mock_request_with_structured_output.num_output_placeholders = 2
result = manager_with_reasoner.should_advance(
mock_request_with_structured_output
)
# placeholder window: start = 5 - 2 = 3, delta = [4, 5]
_, called_delta = reasoner.is_reasoning_end_streaming.call_args[0]
assert list(called_delta) == [4, 5]
assert result is False
def test_should_advance_trims_reasoning_prefix_for_json(
self,
manager_with_reasoner,
mock_request_with_structured_output,
):
"""JSON uses the common trim-then-advance path at the boundary."""
structured_req = mock_request_with_structured_output.structured_output_request
structured_req.reasoning_ended = False
structured_req.structured_output_key = (
StructuredOutputOptions.JSON_OBJECT,
"{}",
)
marker = 248069
class MarkerReasoner:
def __init__(self, *_, **__):
pass
def is_reasoning_end_streaming(self, input_ids, delta_ids):
return marker in list(delta_ids)
structured_req.reasoner = MarkerReasoner()
new_token_ids = [9, 198, marker, 271, 5005]
mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
result = manager_with_reasoner.should_advance(
mock_request_with_structured_output,
new_token_ids=new_token_ids,
)
structured_req.grammar.accept_tokens.assert_not_called()
assert structured_req.reasoning_ended is True
assert result is True
assert structured_req.reasoning_end_token_index == 5
assert manager_with_reasoner.trim_reasoning_for_advance(
mock_request_with_structured_output, new_token_ids
) == [271, 5005]
vllm/models/deepseek_v4/attention.py
+43
View File
@@ -47,6 +47,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.models.utils import extract_layer_index
from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope
from vllm.models.deepseek_v4.compressor import DeepseekCompressor
from vllm.triton_utils import tl, triton
from vllm.utils.multi_stream_utils import (
execute_in_parallel,
maybe_execute_in_parallel,
@@ -66,6 +67,25 @@ from vllm.v1.kv_cache_interface import (
logger = init_logger(__name__)
@triton.jit
def _fill_short_context_topk_indices(
output,
positions,
TOP_K: tl.constexpr,
COMPRESS_RATIO: tl.constexpr,
PADDED_TOP_K: tl.constexpr,
):
# small triton kernel that selects every candidate, -1 otherwise
row = tl.program_id(0)
offsets = tl.arange(0, PADDED_TOP_K)
num_compressed = (tl.load(positions + row) + 1) // COMPRESS_RATIO
tl.store(
output + row * TOP_K + offsets,
tl.where(offsets < num_compressed, offsets, -1),
mask=offsets < TOP_K,
)
def _resolve_dsv4_kv_cache_dtype(
use_fp8_ds_mla_layout: bool,
kv_cache_dtype: str,
@@ -787,6 +807,29 @@ class DeepseekV4Indexer(nn.Module):
) -> torch.Tensor:
compressor = self.compressor
attn_metadata = get_forward_context().attn_metadata
if isinstance(attn_metadata, dict):
indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens:
# candidates num smaller than topk, every candidate is selected
# but we still need to build k cache
compressor(compressed_kv_score, positions, rotary_emb)
assert self.topk_indices_buffer is not None
num_tokens = (
indexer_metadata.num_decode_tokens
+ indexer_metadata.num_prefill_tokens
)
if num_tokens > 0:
_fill_short_context_topk_indices[(num_tokens,)](
self.topk_indices_buffer,
positions,
TOP_K=self.topk_tokens,
COMPRESS_RATIO=self.compress_ratio,
PADDED_TOP_K=triton.next_power_of_2(self.topk_tokens),
num_warps=8,
)
return self.topk_indices_buffer
def wq_b_and_q_quant():
# ReplicatedLinear returns (output, bias); bias is None.
q, _ = self.wq_b(qr)
vllm/v1/core/sched/scheduler.py
+3 -1
View File
@@ -1716,7 +1716,9 @@ class Scheduler(SchedulerInterface):
request.status = RequestStatus.FINISHED_STOPPED
stopped = True
if new_token_ids and self.structured_output_manager.should_advance(request):
if new_token_ids and self.structured_output_manager.should_advance(
request, new_token_ids=new_token_ids
):
struct_output_request = request.structured_output_request
assert struct_output_request is not None
grammar = struct_output_request.grammar
vllm/v1/structured_output/__init__.py
+32 -30
View File
@@ -15,7 +15,6 @@ from vllm.v1.structured_output.backend_guidance import GuidanceBackend
from vllm.v1.structured_output.backend_types import (
StructuredOutputBackend,
StructuredOutputGrammar,
StructuredOutputOptions,
)
from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend
@@ -379,7 +378,11 @@ class StructuredOutputManager:
return request.structured_output_request.reasoning_ended
return True
def should_advance(self, request: "Request") -> bool:
def should_advance(
self,
request: "Request",
new_token_ids: list[int] | None = None,
) -> bool:
if not request.use_structured_output:
return False
@@ -402,37 +405,36 @@ class StructuredOutputManager:
if structured_req.reasoning_ended:
return True
# Check if reasoning ends in *this* step
delta_from = request.num_computed_tokens - request.num_output_placeholders
# Check if reasoning ends in *this* step.
# When the caller passes new_token_ids (the tokens that were just
# appended this step), use it directly as the delta window. The
# placeholder-derived fallback assumes num_output_placeholders ==
# len(new_token_ids), which breaks under async scheduling + spec
# decode when some drafts are rejected (#43388): the placeholder
# count remains > 0 after the step and the computed delta window
# starts past the reasoning-end marker.
all_token_ids = request.all_token_ids
start = (
delta_from if delta_from >= 0 else max(len(all_token_ids) + delta_from, 0)
)
if reasoner.is_reasoning_end_streaming(
all_token_ids, itertools.islice(all_token_ids, start, None)
):
if new_token_ids:
# The tokens were already appended this step, so the step window
# starts exactly len(new_token_ids) from the end.
start = len(all_token_ids) - len(new_token_ids)
delta_ids: Iterable[int] = new_token_ids
else:
delta_from = request.num_computed_tokens - request.num_output_placeholders
start = (
delta_from
if delta_from >= 0
else max(len(all_token_ids) + delta_from, 0)
)
delta_ids = itertools.islice(all_token_ids, start, None)
if reasoner.is_reasoning_end_streaming(all_token_ids, delta_ids):
structured_req.reasoning_ended = True
# Reasoning just ended this step. Defer FSM advance until the next
# pass (see reasoning_ended check above) for JSON/regex/choice/grammar:
# advancing on the closing boundary token can accept tokens that still
# belong to the reasoning stream. Structural tags are the only safe
# same-step exception: they model phased output (e.g. thinking tag ->
# answer tag), and speculative decoding must run grammar.validate_tokens
# on draft tokens produced immediately after that transition.
if (
self.vllm_config.speculative_config is not None
and structured_req.structured_output_key[0]
== StructuredOutputOptions.STRUCTURAL_TAG
):
# The scheduler will advance the grammar with this step's
# tokens right away, but the step still contains reasoning
# content up to and including the end marker. Record where
# it ends so trim_reasoning_for_advance() can drop it.
structured_req.reasoning_end_token_index = (
self._find_reasoning_end_index(reasoner, all_token_ids, start)
)
return True
# Record the boundary so the scheduler can exclude reasoning tokens.
end_index = self._find_reasoning_end_index(reasoner, all_token_ids, start)
structured_req.reasoning_end_token_index = end_index
return True
return False
vllm/v1/worker/gpu/cudagraph_utils.py
+17 -34
View File
@@ -498,7 +498,10 @@ class ModelCudaGraphManager(CudaGraphManager):
block_tables,
attn_groups,
kv_cache_config,
full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
skip_attn=(
desc.cg_mode == CUDAGraphMode.PIECEWISE
and not self.use_breakable_cg
),
)
# Capture with dummy rows marked as padding.
@@ -507,6 +510,7 @@ class ModelCudaGraphManager(CudaGraphManager):
def forward_fn(cg_mode: CUDAGraphMode) -> None:
batch_descriptor = None
if cg_mode == CUDAGraphMode.PIECEWISE:
assert (attn_metadata is not None) == self.use_breakable_cg
batch_descriptor = BatchDescriptor(
num_tokens=num_tokens,
has_lora=has_lora,
@@ -589,7 +593,7 @@ def prepare_inputs_to_capture(
block_tables: BlockTables,
attn_groups: list[list[AttentionGroup]],
kv_cache_config: KVCacheConfig,
full_cudagraph: bool,
skip_attn: bool = False,
) -> AttentionState:
input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
@@ -610,36 +614,15 @@ def prepare_inputs_to_capture(
)
input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs]
# NOTE(woosuk): Attention metadata is required not just by standard attention
# kernels, but also by specialized attention-like operations (e.g., Inkling's sconv,
# DSV4 compressor), which maintain their own states and require special metadata
# such as block tables.
# During CUDA graph capture:
# - For FULL CUDA graphs: We set for_capture=True so that both attention and
# attention-like ops produce capturable metadata compatible with CUDA graphs.
# - For PIECEWISE CUDA graphs: We still build attention metadata, but set
# for_capture=False. This is because:
# * Attention-like ops (such as sconv or DSV4 compressor) may not be used as
# breakpoints in PIECEWISE CUDA graphs, so we must generate their attention
# metadata so they can execute and be captured during graph capture.
# * Standard attention ops that are treated as breakpoints will be executed
# eagerly at capture time (not included in the graph itself), and for these,
# setting for_capture=False is essential. Some attention backends
# (like linear attention) cannot generate capturable metadata for prefill,
# so for_capture=False ensures they execute without issue.
# * We assume that attention-like operations intended for capture will still
# produce capturable metadata, even when for_capture=False. While this
# assumption is brittle, it currently works in practice.
# In summary: We always generate attention metadata for both FULL and PIECEWISE
# CUDA graphs, setting for_capture=True for FULL graphs, and for_capture=False
# for PIECEWISE graphs, to ensure correct execution and capture.
attn_metadata = model_state.prepare_attn(
input_batch,
CUDAGraphMode.NONE,
input_block_tables,
slot_mappings,
attn_groups,
kv_cache_config,
for_capture=full_cudagraph,
)
attn_metadata = None
if not skip_attn:
attn_metadata = model_state.prepare_attn(
input_batch,
CUDAGraphMode.NONE,
input_block_tables,
slot_mappings,
attn_groups,
kv_cache_config,
for_capture=True,
)
return AttentionState(attn_metadata, slot_mappings_by_layer)
vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py
+4 -1
View File
@@ -56,7 +56,10 @@ class SpeculatorCudaGraphManager(CudaGraphManager):
block_tables,
attn_groups,
kv_cache_config,
full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
skip_attn=(
desc.cg_mode == CUDAGraphMode.PIECEWISE
and not self.use_breakable_cg
),
)
return lambda cg_mode: forward_fn(