From 2ad10292339c045db0cb3998a76240a459cc83a7 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Tue, 14 Apr 2026 17:08:17 -0400 Subject: [PATCH 001/150] [Bug] Fix batch invariance nvfp4 support (#39820) Signed-off-by: yewentao256 --- .buildkite/test_areas/misc.yaml | 1 + vllm/model_executor/kernels/linear/__init__.py | 8 +++++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 82fe6670a99..0cf9ec43392 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -224,6 +224,7 @@ steps: - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA] - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN] + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py - label: Acceptance Length Test (Large Models) # optional timeout_in_minutes: 25 diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index d0c0dd9f328..e0a9272c890 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -601,7 +601,13 @@ def init_nvfp4_linear_kernel() -> NvFp4LinearKernel: # Env-var overrides. force_kernel: type[NvFp4LinearKernel] | None = None - if envs.VLLM_USE_FBGEMM: + if envs.VLLM_BATCH_INVARIANT: + logger.info_once( + "VLLM_BATCH_INVARIANT forces NVFP4 linear to use the " + "emulation backend for deterministic execution." + ) + force_kernel = EmulationNvFp4LinearKernel + elif envs.VLLM_USE_FBGEMM: force_kernel = FbgemmNvFp4LinearKernel elif envs.VLLM_USE_NVFP4_CT_EMULATIONS: force_kernel = EmulationNvFp4LinearKernel From 1696c864b90afb9cec925312d0cf12c396517007 Mon Sep 17 00:00:00 2001 From: Zhewen Li Date: Tue, 14 Apr 2026 14:13:58 -0700 Subject: [PATCH 002/150] [Bugfix][Mooncake] Fix thread-local CUDA context for NVLink transfers in _send_blocks (#39548) Signed-off-by: Zhewen Li Co-authored-by: Zhewen Li --- .../kv_connector/v1/mooncake/mooncake_connector.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index b49a016641e..6c17bf82ace 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -41,6 +41,7 @@ from vllm.distributed.parallel_state import ( ) from vllm.forward_context import ForwardContext from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.utils import get_kv_cache_layout @@ -645,6 +646,10 @@ class MooncakeConnectorWorker: logger.info("Initializing Mooncake Transfer Engine worker %s", engine_id) self.vllm_config = vllm_config + # Capture device BEFORE TransferEngine init — MNNVL's NVLink allocator + # may change the current CUDA device during engine.initialize(). + self.device_id = torch.accelerator.current_device_index() + current_platform.set_device(self.device_id) self.engine = TransferEngine() self.hostname = get_ip() @@ -705,9 +710,12 @@ class MooncakeConnectorWorker: # For kv_both, we will act both prefiller and decoder. if not self.is_kv_consumer: # Background threads for sending kvcaches to D. + # Each pool thread must be bound to the correct CUDA device + # because CUDA device selection is thread-local. self._sender_executor = ThreadPoolExecutor( max_workers=self.num_sender_workers, thread_name_prefix="vllm-mooncake-sender", + initializer=self._bind_sender_thread_device, ) logger.debug( "Mooncake Prefiller: use %d workers to send kvcaches", @@ -1193,6 +1201,12 @@ class MooncakeConnectorWorker: return src_ptrs, dst_ptrs, lengths, err_reqs, err_msg + def _bind_sender_thread_device(self) -> None: + """ThreadPoolExecutor initializer — binds each pool thread to the + correct CUDA device. CUDA device selection is thread-local, so + without this, NVLink transfers fail for TP ranks > 0.""" + current_platform.set_device(self.device_id) + def _send_blocks( self, remote_session: str, From 507df79a29ac8733ea03c624d215ccfffad7bbe9 Mon Sep 17 00:00:00 2001 From: Francesco Fusco Date: Wed, 15 Apr 2026 00:19:09 +0200 Subject: [PATCH 003/150] [Hybrid] Simplify accepted token counting in spec decode for hybrid models (#38372) --- vllm/v1/worker/gpu_model_runner.py | 23 ++++------------------- 1 file changed, 4 insertions(+), 19 deletions(-) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e8df720e78c..4c573f79e97 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1428,26 +1428,11 @@ class GPUModelRunner( # TODO: Remove .cpu() sync to enable fully async for hybrid model; # Use num_computed_tokens.gpu instead of req.num_computed_tokens to # support aligned mamba cache mode. - # Find the number of accepted tokens for each sequence. + # Count the number of accepted tokens for each sequence. + # Valid tokens are contiguous from position 0, so counting non-(-1) + # tokens gives us the first -1 position (i.e., number of accepted). num_reqs = output_token_ids.size(0) - self.num_accepted_tokens.gpu[:num_reqs] = ( - ( - torch.cat( - [ - output_token_ids, - torch.full( - (num_reqs, 1), - -1, - device=output_token_ids.device, - ), - ], - dim=1, - ) - == -1 - ) - .int() - .argmax(-1) - ) + self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum(dim=1) if self.cache_config.mamba_cache_mode == "align": for i, num_tokens in enumerate( From 65b9808960d7f25a018c525d2aca3dba7c411cfb Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Tue, 14 Apr 2026 19:03:57 -0400 Subject: [PATCH 004/150] [Bugfix] Disable FlashInfer CUTLASS MoE on SM121 (DGX Spark) (#39825) Signed-off-by: mgoin Co-authored-by: Claude --- .../layers/fused_moe/flashinfer_cutlass_moe.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py index 26409804c48..0d47b0f3174 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py @@ -130,7 +130,14 @@ class FlashInferExperts(mk.FusedMoEExpertsModular): p.is_device_capability(90) or p.is_device_capability_family(100) or p.is_device_capability_family(110) - or p.is_device_capability_family(120) + or p.is_device_capability(120) + # NOTE: SM121 (DGX Spark) is excluded because the bf16 + # unquantized CUTLASS MoE GEMM in flashinfer <= 0.6.7 has no + # Relu2 template instantiation and throws "Invalid activation + # type" on Nemotron-H. Fixed upstream by + # https://github.com/flashinfer-ai/flashinfer/pull/2926 + # (merged 2026-04-01, not yet in a stable release); lift this + # restriction once flashinfer >= 0.6.8 is the minimum. ) and has_flashinfer_cutlass_fused_moe() ) From b569620f720d0c07bcb90bcc29dacaf0ce318d85 Mon Sep 17 00:00:00 2001 From: Andrey Talman Date: Tue, 14 Apr 2026 20:13:24 -0400 Subject: [PATCH 005/150] [CI] Add PyTorch nightly build and test pipeline (#37226) Signed-off-by: atalman --- .../image_build/image_build_torch_nightly.sh | 68 +++++++++++++++++++ .buildkite/test_areas/models_basic.yaml | 11 +-- .buildkite/test_areas/models_language.yaml | 13 ++-- 3 files changed, 83 insertions(+), 9 deletions(-) create mode 100755 .buildkite/image_build/image_build_torch_nightly.sh diff --git a/.buildkite/image_build/image_build_torch_nightly.sh b/.buildkite/image_build/image_build_torch_nightly.sh new file mode 100755 index 00000000000..a23c658d46b --- /dev/null +++ b/.buildkite/image_build/image_build_torch_nightly.sh @@ -0,0 +1,68 @@ +#!/bin/bash +set -euo pipefail + +# Build a vLLM test image with PyTorch nightly installed. +# Called by the pipeline generator's "vLLM Against PyTorch Nightly" group. + +if [[ $# -lt 5 ]]; then + echo "Usage: $0 " + exit 1 +fi + +REGISTRY=$1 +REPO=$2 +BUILDKITE_COMMIT=$3 +BRANCH=$4 +IMAGE_TAG=$5 + +# --- Arguments --- +echo "--- :mag: Arguments" +echo "REGISTRY: ${REGISTRY}" +echo "REPO: ${REPO}" +echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}" +echo "BRANCH: ${BRANCH}" +echo "IMAGE_TAG: ${IMAGE_TAG}" + +# --- ECR login --- +echo "--- :key: ECR login" +aws ecr-public get-login-password --region us-east-1 \ + | docker login --username AWS --password-stdin "$REGISTRY" +aws ecr get-login-password --region us-east-1 \ + | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com + +# --- Set up buildx --- +echo "--- :docker: Setting up buildx" +docker buildx create --name vllm-builder --driver docker-container --use || true +docker buildx inspect --bootstrap +docker buildx ls + +# --- Skip if image already exists --- +echo "--- :mag: Checking if image already exists" +if docker manifest inspect "$IMAGE_TAG" >/dev/null 2>&1; then + echo "Image found: $IMAGE_TAG — skipping build" + exit 0 +fi +echo "Image not found, proceeding with build..." + +# --- CUDA 13.0 for nightly builds --- +# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9 +NIGHTLY_CUDA_VERSION="13.0.0" +NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04" +NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04" + +echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})" +docker buildx build --file docker/Dockerfile \ + --build-arg max_jobs=16 \ + --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ + --build-arg USE_SCCACHE=1 \ + --build-arg PYTORCH_NIGHTLY=1 \ + --build-arg CUDA_VERSION="${NIGHTLY_CUDA_VERSION}" \ + --build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \ + --build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \ + --build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \ + --tag "$IMAGE_TAG" \ + --push \ + --target test \ + --progress plain . + +echo "--- :white_check_mark: Torch nightly image build complete: $IMAGE_TAG" diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 8ba4484fb02..10b038d8b8a 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -1,5 +1,5 @@ group: Models - Basic -depends_on: +depends_on: - image-build steps: - label: Basic Models Tests (Initialization) @@ -13,10 +13,11 @@ steps: commands: # Run a subset of model initialization tests - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset + mirror: + torch_nightly: {} - label: Basic Models Tests (Extra Initialization) %N timeout_in_minutes: 45 - torch_nightly: true source_file_dependencies: - vllm/model_executor/models/ - tests/models/test_initialization.py @@ -27,6 +28,8 @@ steps: # test.) Also run if model initialization test file is modified - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 2 + mirror: + torch_nightly: {} - label: Basic Models Tests (Other) timeout_in_minutes: 45 @@ -42,10 +45,10 @@ steps: device: mi325_1 depends_on: - image-build-amd - + - label: Basic Models Test (Other CPU) # 5min - depends_on: + depends_on: - image-build-cpu timeout_in_minutes: 10 source_file_dependencies: diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 7eac9e30193..c13371e25f1 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -1,10 +1,9 @@ group: Models - Language -depends_on: +depends_on: - image-build steps: - label: Language Models Tests (Standard) timeout_in_minutes: 25 - torch_nightly: true source_file_dependencies: - vllm/ - tests/models/language @@ -12,10 +11,11 @@ steps: # Test standard language models, excluding a subset of slow tests - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and (not slow_test)' + mirror: + torch_nightly: {} - label: Language Models Tests (Extra Standard) %N timeout_in_minutes: 45 - torch_nightly: true source_file_dependencies: - vllm/model_executor/models/ - tests/models/language/pooling/test_embedding.py @@ -27,10 +27,11 @@ steps: - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 2 + mirror: + torch_nightly: {} - label: Language Models Tests (Hybrid) %N timeout_in_minutes: 75 - torch_nightly: true source_file_dependencies: - vllm/ - tests/models/language/generation @@ -42,6 +43,8 @@ steps: # Shard hybrid language model tests - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 2 + mirror: + torch_nightly: {} - label: Language Models Test (Extended Generation) # 80min timeout_in_minutes: 110 @@ -62,7 +65,7 @@ steps: - image-build-amd commands: - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2' + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' - label: Language Models Test (PPL) From 3bfe55a03758d57d4dde7db975842dc281b740a0 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Tue, 14 Apr 2026 17:47:57 -0700 Subject: [PATCH 006/150] [Model Runner V2] Disable piecewise cudagraph mode fallback for eagle draft decodes (#39773) Signed-off-by: Giancarlo Delfin --- .../gpu/spec_decode/eagle/speculator.py | 28 +++++++++++++------ 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 881858f5d24..b1126983539 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -115,13 +115,21 @@ class EagleSpeculator: cudagraph_mode, self.num_speculative_steps + 1, ) - # Initialize cudagraph manager for draft generation (draft positions > 0). + + # PIECEWISE cudagraphs are not supported for eagle draft decodes. + # PIECEWISE pads num_tokens to the next capture size without padding + # num_reqs, which can cause attention backends to read past the + # valid per-request metadata (e.g. FlashInfer's kv_indptr buffer). + if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL: + cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + else: + cudagraph_mode = CUDAGraphMode.NONE + + # Initialize cudagraph manager for draft decodes (draft positions > 0). self.decode_cudagraph_manager = EagleCudaGraphManager( self.vllm_config, self.device, - # Only use FULL graph mode, if available, because draft decodes - # only consist of a single token. - cudagraph_mode.decode_mode(), + cudagraph_mode, decode_query_len=1, ) # Share a single pool between prefill and decode since they never @@ -366,11 +374,8 @@ class EagleSpeculator: # 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. For PIECEWISE, only the model's compiled regions - # are captured, and the rest (compute_logits, gumbel_sample, - # update_eagle_inputs) runs eagerly. + # each step). For FULL graphs, the entire multi-step loop is + # recorded as one graph. assert self.decode_cudagraph_manager is not None self.decode_cudagraph_manager.capture( self.generate_draft, @@ -629,6 +634,11 @@ def _prepare_eagle_inputs_kernel( block = i + tl.arange(0, BLOCK_SIZE) mask = block < max_num_reqs tl.store(eagle_seq_lens_ptr + block, 0, mask=mask) + # Pad last_token_indices for CUDA graphs. + for i in range(num_reqs, max_num_reqs, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < max_num_reqs + tl.store(last_token_indices_ptr + block, 0, mask=mask) def prepare_eagle_inputs( From f4b42df04847dcbd3247f7f4c56dff45e40bdf0d Mon Sep 17 00:00:00 2001 From: Vibhav Agarwal Date: Wed, 15 Apr 2026 08:27:13 +0530 Subject: [PATCH 007/150] [Attention Backend] TurboQuant: 2-bit KV cache compression with 4x capacity (#38479) Signed-off-by: vibhavagarwal5 Signed-off-by: Michael Goin Co-authored-by: Xinyu Chen Co-authored-by: Michael Goin --- .buildkite/test_areas/lm_eval.yaml | 10 + docs/design/attention_backends.md | 1 + pyproject.toml | 3 + .../gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml | 5 + .../evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml | 5 + .../evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml | 5 + .../evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml | 5 + .../evals/gsm8k/configs/models-turboquant.txt | 4 + tests/quantization/test_turboquant.py | 570 ++++++++++++ vllm/config/attention.py | 5 + vllm/config/cache.py | 4 + vllm/engine/arg_utils.py | 38 + .../layers/attention/attention.py | 82 ++ .../quantization/turboquant/__init__.py | 14 + .../quantization/turboquant/centroids.py | 86 ++ .../layers/quantization/turboquant/config.py | 185 ++++ .../quantization/turboquant/quantizer.py | 24 + vllm/platforms/cuda.py | 5 + vllm/platforms/xpu.py | 6 + vllm/utils/torch_utils.py | 4 + vllm/v1/attention/backends/registry.py | 1 + vllm/v1/attention/backends/turboquant_attn.py | 812 ++++++++++++++++++ .../attention/ops/triton_turboquant_decode.py | 617 +++++++++++++ .../attention/ops/triton_turboquant_store.py | 441 ++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 6 +- vllm/v1/kv_cache_interface.py | 26 + vllm/v1/worker/utils.py | 2 +- 27 files changed, 2963 insertions(+), 3 deletions(-) create mode 100644 tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml create mode 100644 tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml create mode 100644 tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml create mode 100644 tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml create mode 100644 tests/evals/gsm8k/configs/models-turboquant.txt create mode 100644 tests/quantization/test_turboquant.py create mode 100644 vllm/model_executor/layers/quantization/turboquant/__init__.py create mode 100644 vllm/model_executor/layers/quantization/turboquant/centroids.py create mode 100644 vllm/model_executor/layers/quantization/turboquant/config.py create mode 100644 vllm/model_executor/layers/quantization/turboquant/quantizer.py create mode 100644 vllm/v1/attention/backends/turboquant_attn.py create mode 100644 vllm/v1/attention/ops/triton_turboquant_decode.py create mode 100644 vllm/v1/attention/ops/triton_turboquant_store.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 39029efe9cd..a07d702cf3c 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -91,6 +91,16 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt +- label: LM Eval TurboQuant KV Cache + timeout_in_minutes: 75 + source_file_dependencies: + - vllm/model_executor/layers/quantization/turboquant/ + - vllm/v1/attention/backends/turboquant_attn.py + - vllm/v1/attention/ops/triton_turboquant_decode.py + - vllm/v1/attention/ops/triton_turboquant_store.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt + - label: GPQA Eval (GPT-OSS) (H100) timeout_in_minutes: 120 device: h100 diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 242cc6b3b1e..65e93657e78 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -178,6 +178,7 @@ Priority is **1 = highest** (tried first). | `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`. > diff --git a/pyproject.toml b/pyproject.toml index c4951b1a4ed..f55dd9308bd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -170,6 +170,9 @@ eles = "eles" datas = "datas" ser = "ser" ure = "ure" +# Walsh-Hadamard Transform +wht = "wht" +WHT = "WHT" [tool.uv] no-build-isolation-package = ["torch"] diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml new file mode 100644 index 00000000000..fedb7416960 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3-4B" +accuracy_threshold: 0.78 +num_questions: 1319 +num_fewshot: 5 +server_args: "--kv-cache-dtype turboquant_k3v4_nc --enforce-eager --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml new file mode 100644 index 00000000000..9717333582b --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3-4B" +accuracy_threshold: 0.80 +num_questions: 1319 +num_fewshot: 5 +server_args: "--kv-cache-dtype turboquant_k8v4 --enforce-eager --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml new file mode 100644 index 00000000000..8ece1852625 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3-4B" +accuracy_threshold: 0.75 +num_questions: 1319 +num_fewshot: 5 +server_args: "--kv-cache-dtype turboquant_3bit_nc --enforce-eager --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml new file mode 100644 index 00000000000..9b3a14f9b95 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3-4B" +accuracy_threshold: 0.80 +num_questions: 1319 +num_fewshot: 5 +server_args: "--kv-cache-dtype turboquant_4bit_nc --enforce-eager --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/models-turboquant.txt b/tests/evals/gsm8k/configs/models-turboquant.txt new file mode 100644 index 00000000000..518aac780b9 --- /dev/null +++ b/tests/evals/gsm8k/configs/models-turboquant.txt @@ -0,0 +1,4 @@ +Qwen3-4B-TQ-k8v4.yaml +Qwen3-4B-TQ-t4nc.yaml +Qwen3-4B-TQ-k3v4nc.yaml +Qwen3-4B-TQ-t3nc.yaml diff --git a/tests/quantization/test_turboquant.py b/tests/quantization/test_turboquant.py new file mode 100644 index 00000000000..78c137e6762 --- /dev/null +++ b/tests/quantization/test_turboquant.py @@ -0,0 +1,570 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for TurboQuant KV-cache quantization. + +Run: .venv/bin/python -m pytest tests/quantization/test_turboquant.py -v +""" + +import math + +import pytest +import torch + +from vllm.model_executor.layers.quantization.turboquant.centroids import ( + get_centroids, + solve_lloyd_max, +) +from vllm.model_executor.layers.quantization.turboquant.config import ( + TQ_PRESETS, + TurboQuantConfig, +) +from vllm.model_executor.layers.quantization.turboquant.quantizer import ( + generate_wht_signs, +) +from vllm.utils.math_utils import next_power_of_2 + +# ============================================================================ +# Helpers +# ============================================================================ + +ALL_PRESETS = list(TQ_PRESETS.keys()) + + +def _assert_strictly_sorted(seq, name="sequence"): + for i in range(len(seq) - 1): + assert seq[i] < seq[i + 1], f"{name} not sorted at index {i}" + + +def _is_power_of_2(n: int) -> bool: + return n > 0 and next_power_of_2(n) == n + + +# Expected concrete values for each preset at head_dim=128. +# fmt: off +PRESET_EXPECTED = { + "turboquant_k8v4": dict( + key_fp8=True, key_quant_bits=8, + key_mse_bits=0, value_quant_bits=4, + mse_bits=4, n_centroids=16, centroid_bits=4, + norm_correction=False, + key_packed_size=128, value_packed_size=68, + slot_size=196, slot_size_aligned=196, + ), + "turboquant_4bit_nc": dict( + key_fp8=False, key_quant_bits=4, + key_mse_bits=4, value_quant_bits=4, + mse_bits=4, n_centroids=16, centroid_bits=4, + norm_correction=True, + key_packed_size=66, value_packed_size=68, + slot_size=134, slot_size_aligned=134, + ), + "turboquant_k3v4_nc": dict( + key_fp8=False, key_quant_bits=3, + key_mse_bits=3, value_quant_bits=4, + mse_bits=3, n_centroids=8, centroid_bits=3, + norm_correction=True, + key_packed_size=50, value_packed_size=68, + slot_size=118, slot_size_aligned=118, + ), + "turboquant_3bit_nc": dict( + key_fp8=False, key_quant_bits=3, + key_mse_bits=3, value_quant_bits=3, + mse_bits=3, n_centroids=8, centroid_bits=3, + norm_correction=True, + key_packed_size=50, value_packed_size=52, + slot_size=102, slot_size_aligned=102, + ), +} +# fmt: on + + +# ============================================================================ +# Config tests (CPU-only, no dependencies beyond config.py) +# ============================================================================ + + +class TestTurboQuantConfig: + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_preset_parses(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert isinstance(cfg, TurboQuantConfig) + + def test_invalid_preset_raises(self): + with pytest.raises(ValueError, match="Unknown TurboQuant"): + TurboQuantConfig.from_cache_dtype("turboquant_invalid", head_dim=128) + + # ---- Per-preset concrete value checks (table-driven) ---- + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_key_mode(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + exp = PRESET_EXPECTED[preset] + assert cfg.key_fp8 is exp["key_fp8"] + assert cfg.key_quant_bits == exp["key_quant_bits"] + assert cfg.key_mse_bits == exp["key_mse_bits"] + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_value_mode(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + exp = PRESET_EXPECTED[preset] + assert cfg.value_quant_bits == exp["value_quant_bits"] + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_bits_and_centroids(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + exp = PRESET_EXPECTED[preset] + assert cfg.mse_bits == exp["mse_bits"] + assert cfg.n_centroids == exp["n_centroids"] + assert cfg.centroid_bits == exp["centroid_bits"] + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_norm_correction(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.norm_correction is PRESET_EXPECTED[preset]["norm_correction"] + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_packed_sizes(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + exp = PRESET_EXPECTED[preset] + assert cfg.key_packed_size == exp["key_packed_size"] + assert cfg.value_packed_size == exp["value_packed_size"] + assert cfg.slot_size == exp["slot_size"] + assert cfg.slot_size_aligned == exp["slot_size_aligned"] + + # ---- Cross-preset structural invariants ---- + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_slot_equals_key_plus_value(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_padded_slot_is_even(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.slot_size_aligned >= cfg.slot_size + assert cfg.slot_size_aligned % 2 == 0, ( + f"slot_size_aligned={cfg.slot_size_aligned} is not even" + ) + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_key_value_packed_sizes_positive(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.key_packed_size > 0 + assert cfg.value_packed_size > 0 + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_n_centroids_is_2_to_mse_bits(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.n_centroids == 2**cfg.mse_bits + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_centroid_bits_always_positive(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.centroid_bits > 0 + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_mse_key_or_fp8_exclusive(self, preset): + """Each preset is either FP8 keys or MSE keys, never both.""" + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + if cfg.key_fp8: + assert cfg.key_mse_bits == 0 + assert cfg.key_quant_bits == 8 + else: + assert cfg.key_mse_bits > 0 + assert cfg.key_quant_bits in (3, 4) + + @pytest.mark.parametrize("preset", ALL_PRESETS) + @pytest.mark.parametrize("head_dim", [64, 96, 128, 256]) + def test_all_presets_all_head_dims(self, preset, head_dim): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=head_dim) + assert cfg.head_dim == head_dim + assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size + assert cfg.slot_size_aligned >= cfg.slot_size + assert cfg.slot_size_aligned % 2 == 0 + + # ---- Boundary skip layers ---- + + def test_boundary_skip_layers_basic(self): + layers = TurboQuantConfig.get_boundary_skip_layers(32) + assert layers == ["0", "1", "30", "31"] + + def test_boundary_skip_layers_zero(self): + assert TurboQuantConfig.get_boundary_skip_layers(32, 0) == [] + + def test_boundary_skip_layers_small_model(self): + layers = TurboQuantConfig.get_boundary_skip_layers(4) + assert layers == ["0", "1", "2", "3"] + + def test_boundary_skip_layers_cap_at_half(self): + layers = TurboQuantConfig.get_boundary_skip_layers(8, 10) + assert len(layers) == 8 + + +# ============================================================================ +# Centroids tests (CPU-only) +# ============================================================================ + + +class TestCentroids: + @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)]) + def test_centroids_shape(self, bits, expected_n): + c = get_centroids(128, bits) + assert c.shape == (expected_n,) + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_centroids_sorted(self, bits): + _assert_strictly_sorted(get_centroids(128, bits), "centroids") + + def test_centroids_cached(self): + c1 = get_centroids(128, 3) + c2 = get_centroids(128, 3) + assert c1 is c2, "get_centroids should return cached object" + + def test_centroids_different_dims_not_identical(self): + c64 = get_centroids(64, 3) + c128 = get_centroids(128, 3) + assert not torch.equal(c64, c128) + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_centroids_symmetric_around_zero(self, bits): + """N(0, 1/d) is symmetric, so centroids should be ~symmetric.""" + c = get_centroids(128, bits) + assert abs(c.mean().item()) < 0.01, "Centroids not centered near 0" + assert abs(c[0].item() + c[-1].item()) < 0.01 + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_centroids_within_4sigma(self, bits): + """All centroids should be within ~4 sigma of N(0, 1/d).""" + sigma = math.sqrt(1.0 / 128) + c = get_centroids(128, bits) + for i, val in enumerate(c): + assert abs(val.item()) < 4 * sigma, ( + f"Centroid {i}={val:.6f} outside 4*sigma={4 * sigma:.6f}" + ) + + +class TestLloydMax: + @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)]) + def test_solve_shapes(self, bits, expected_n): + centroids, boundaries = solve_lloyd_max(128, bits) + assert centroids.shape == (expected_n,) + assert boundaries.shape == (expected_n - 1,) + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_centroids_sorted(self, bits): + centroids, _ = solve_lloyd_max(128, bits) + _assert_strictly_sorted(centroids, "centroids") + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_boundaries_sorted(self, bits): + _, boundaries = solve_lloyd_max(128, bits) + _assert_strictly_sorted(boundaries, "boundaries") + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_boundaries_between_centroids(self, bits): + """Each boundary must lie between its adjacent centroids.""" + centroids, boundaries = solve_lloyd_max(128, bits) + for i in range(len(boundaries)): + assert centroids[i] < boundaries[i] < centroids[i + 1], ( + f"Boundary {i}={boundaries[i]:.6f} not between " + f"c[{i}]={centroids[i]:.6f} and c[{i + 1}]={centroids[i + 1]:.6f}" + ) + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_boundaries_are_midpoints(self, bits): + """Lloyd-Max boundaries are midpoints of adjacent centroids.""" + centroids, boundaries = solve_lloyd_max(128, bits) + for i in range(len(boundaries)): + expected = (centroids[i] + centroids[i + 1]) / 2.0 + assert abs(boundaries[i].item() - expected.item()) < 1e-6 + + def test_solve_deterministic(self): + c1, b1 = solve_lloyd_max(128, 3) + c2, b2 = solve_lloyd_max(128, 3) + assert torch.equal(c1, c2) + assert torch.equal(b1, b2) + + def test_solve_dtype_float32(self): + centroids, boundaries = solve_lloyd_max(128, 3) + assert centroids.dtype == torch.float32 + assert boundaries.dtype == torch.float32 + + @pytest.mark.parametrize("bits", [3, 4]) + def test_centroids_match_scipy_reference(self, bits): + """Verify _trapz(n=200) centroids match scipy.integrate.quad reference. + + This ensures our scipy-free trapezoid integration doesn't silently + drift from the published Lloyd-Max quality. + """ + pytest.importorskip("scipy") + from scipy.integrate import quad + + d = 128 + sigma2 = 1.0 / d + sigma = math.sqrt(sigma2) + + def pdf(x): + return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp( + -x * x / (2 * sigma2) + ) + + n_levels = 2**bits + lo, hi = -3.5 * sigma, 3.5 * sigma + ref_centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)] + for _ in range(200): + boundaries = [ + (ref_centroids[i] + ref_centroids[i + 1]) / 2.0 + for i in range(n_levels - 1) + ] + edges = [lo * 3] + boundaries + [hi * 3] + new_centroids = [] + for i in range(n_levels): + a, b = edges[i], edges[i + 1] + num, _ = quad(lambda x: x * pdf(x), a, b) + den, _ = quad(pdf, a, b) + new_centroids.append(num / den if den > 1e-15 else ref_centroids[i]) + if ( + max(abs(new_centroids[i] - ref_centroids[i]) for i in range(n_levels)) + < 1e-10 + ): + break + ref_centroids = new_centroids + + # Compare our _trapz centroids against scipy reference + our_centroids, _ = solve_lloyd_max(d, bits) + ref_t = torch.tensor(ref_centroids, dtype=torch.float32) + max_err = (our_centroids - ref_t).abs().max().item() + # _trapz(n=200) has ~O(h^2) error vs adaptive quad; 1e-3 is tight + # enough to catch regression while allowing trapezoid approximation. + assert max_err < 1e-3, ( + f"d={d}, bits={bits}: max centroid error vs scipy = {max_err:.2e}" + ) + + +# ============================================================================ +# Rotation matrix tests (GPU required) +# ============================================================================ + +CUDA_AVAILABLE = torch.cuda.is_available() + + +def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor: + """Haar-distributed random orthogonal matrix via QR (test/benchmark only).""" + gen = torch.Generator(device="cpu") + gen.manual_seed(seed) + G = torch.randn(d, d, generator=gen, device="cpu", dtype=torch.float32) + Q, R = torch.linalg.qr(G) + diag_sign = torch.sign(torch.diag(R)) + diag_sign[diag_sign == 0] = 1.0 + Q = Q * diag_sign.unsqueeze(0) + return Q.to(device) + + +@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available") +class TestRotationMatrix: + """Tests for the QR-based rotation (standalone benchmarks only).""" + + @pytest.mark.parametrize("dim", [64, 96, 128, 256]) + def test_rotation_matrix_shape_and_orthogonal(self, dim): + Pi = generate_rotation_matrix(dim, seed=42, device="cuda") + assert Pi.shape == (dim, dim) + eye = Pi @ Pi.T + assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), ( + f"Pi not orthogonal for dim={dim}" + ) + + def test_rotation_matrix_deterministic(self): + Pi1 = generate_rotation_matrix(128, seed=42) + Pi2 = generate_rotation_matrix(128, seed=42) + assert torch.equal(Pi1, Pi2) + + def test_rotation_matrix_different_seeds(self): + Pi1 = generate_rotation_matrix(128, seed=42) + Pi2 = generate_rotation_matrix(128, seed=99) + assert not torch.equal(Pi1, Pi2) + + def test_rotation_matrix_det_is_pm1(self): + """Orthogonal matrix determinant must be +1 or -1.""" + Pi = generate_rotation_matrix(128, seed=42, device="cuda") + det = torch.linalg.det(Pi) + assert abs(abs(det.item()) - 1.0) < 1e-4 + + +# ============================================================================ +# WHT rotation tests (serving path: generate_wht_signs + _build_hadamard) +# ============================================================================ + + +def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor: + """Reproduce the serving-path Hadamard construction.""" + H = torch.tensor([[1.0]]) + while H.shape[0] < d: + H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0) + return (H / math.sqrt(d)).to(torch.device(device)) + + +@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available") +class TestWHTRotation: + """Tests for the WHT rotation actually used in serving.""" + + @pytest.mark.parametrize("dim", [64, 128, 256]) + def test_wht_orthonormal(self, dim): + """signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I.""" + signs = generate_wht_signs(dim, seed=42, device="cuda") + H = _build_hadamard(dim, "cuda") + PiT = (signs.unsqueeze(1) * H).contiguous() + eye = PiT @ PiT.T + assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), ( + f"WHT rotation not orthonormal for dim={dim}" + ) + + @pytest.mark.parametrize("dim", [64, 128, 256]) + def test_wht_self_inverse(self, dim): + """PiT should be self-inverse: PiT @ PiT = I (up to sign flip).""" + signs = generate_wht_signs(dim, seed=42, device="cuda") + H = _build_hadamard(dim, "cuda") + PiT = (signs.unsqueeze(1) * H).contiguous() + Pi = PiT.T.contiguous() + # Pi @ PiT should be identity (rotation then inverse) + result = Pi @ PiT + assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), ( + f"WHT rotation not self-inverse for dim={dim}" + ) + + def test_wht_signs_deterministic(self): + """Same seed must produce identical signs.""" + s1 = generate_wht_signs(128, seed=42) + s2 = generate_wht_signs(128, seed=42) + assert torch.equal(s1, s2) + + def test_wht_signs_different_seeds(self): + """Different seeds must produce different signs.""" + s1 = generate_wht_signs(128, seed=42) + s2 = generate_wht_signs(128, seed=99) + assert not torch.equal(s1, s2) + + def test_wht_signs_are_pm1(self): + """All sign values must be exactly +1 or -1.""" + signs = generate_wht_signs(128, seed=42) + assert torch.all(signs.abs() == 1.0) + + +# ============================================================================ +# Store → Decode round-trip test (GPU + Triton required) +# ============================================================================ + + +@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available") +class TestStoreDecodeRoundTrip: + """End-to-end: store KV into TQ cache, decode, compare vs fp16 ref.""" + + @pytest.mark.parametrize( + "preset", + ["turboquant_k8v4", "turboquant_4bit_nc"], + ) + def test_single_token_roundtrip(self, preset): + """Store 1 token, decode with query=key, check attention output. + + For a single token with query=key, attention output should equal + the value (softmax over single key = 1.0). Quantization error + means we check cosine similarity rather than exact equality. + """ + from vllm.model_executor.layers.quantization.turboquant.centroids import ( + solve_lloyd_max, + ) + from vllm.v1.attention.ops.triton_turboquant_decode import ( + triton_turboquant_decode_attention, + ) + from vllm.v1.attention.ops.triton_turboquant_store import ( + triton_turboquant_store, + ) + + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + D = 128 + Hk = 4 # num_kv_heads + Hq = 4 # num_q_heads (no GQA for simplicity) + B = 1 # single token + block_size = 16 + num_blocks = 1 + + device = torch.device("cuda") + + # Generate rotation + signs = generate_wht_signs(D, seed=42, device=device) + H = _build_hadamard(D, "cuda") + PiT = (signs.unsqueeze(1) * H).contiguous().float() + Pi = PiT.T.contiguous() + + # Generate centroids + centroids, _ = solve_lloyd_max(D, cfg.centroid_bits) + centroids = centroids.float().to(device) + c_sorted, _ = centroids.sort() + midpoints = ((c_sorted[:-1] + c_sorted[1:]) / 2).to(device) + + # Random K, V + torch.manual_seed(123) + key = torch.randn(B, Hk, D, device=device, dtype=torch.float16) + value = torch.randn(B, Hk, D, device=device, dtype=torch.float16) + + # Allocate KV cache + padded_slot = cfg.slot_size_aligned + kv_cache = torch.zeros( + num_blocks, + block_size, + Hk, + padded_slot, + device=device, + dtype=torch.uint8, + ) + slot_mapping = torch.tensor([0], device=device, dtype=torch.int32) + + # Store + triton_turboquant_store( + key, + value, + kv_cache, + slot_mapping, + PiT, + midpoints, + mse_bits=cfg.key_mse_bits, + key_packed_size=cfg.key_packed_size, + value_quant_bits=cfg.effective_value_quant_bits, + key_fp8=cfg.key_fp8, + ) + + # Decode: use key as query so attention = softmax([1]) * V = V + query = key.expand(B, Hq, D).contiguous().to(torch.float16) + block_table = torch.tensor([[0]], device=device, dtype=torch.int32) + seq_lens = torch.tensor([1], device=device, dtype=torch.int32) + + output = triton_turboquant_decode_attention( + query=query, + kv_cache=kv_cache, + block_table=block_table, + seq_lens=seq_lens, + Pi=Pi, + centroids=centroids, + scale=1.0 / math.sqrt(D), + mse_bits=cfg.key_mse_bits, + key_packed_size=cfg.key_packed_size, + value_quant_bits=cfg.effective_value_quant_bits, + key_fp8=cfg.key_fp8, + norm_correction=cfg.norm_correction, + PiT=PiT, + max_num_kv_splits=4, + ) + + # With single KV, output should approximate the stored value. + # Check per-head cosine similarity > threshold. + out_fp32 = output.float() + val_fp32 = value.expand(B, Hq, D).float() + for h in range(Hq): + cos_sim = torch.nn.functional.cosine_similarity( + out_fp32[0, h].unsqueeze(0), + val_fp32[0, h].unsqueeze(0), + ).item() + # FP8 keys should be very accurate; MSE keys have more error + threshold = 0.95 if cfg.key_fp8 else 0.85 + assert cos_sim > threshold, ( + f"Preset {preset} head {h}: cosine_sim={cos_sim:.4f} < {threshold}" + ) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 1da647a6d6f..561367173d5 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -27,6 +27,11 @@ class AttentionConfig: flash_attn_max_num_splits_for_cuda_graph: int = 32 """Flash Attention max number splits for cuda graph decode.""" + tq_max_kv_splits_for_cuda_graph: int = 32 + """TurboQuant max NUM_KV_SPLITS for cuda graph decode. + Fixes the split count so grid dimensions are constant across captures, + and buffers can be pre-allocated to avoid inflating the memory estimate.""" + use_cudnn_prefill: bool = False """Whether to use cudnn prefill.""" diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 20721cc8092..47a655f22d5 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -24,6 +24,10 @@ CacheDType = Literal[ "fp8_e5m2", "fp8_inc", "fp8_ds_mla", + "turboquant_k8v4", + "turboquant_4bit_nc", + "turboquant_k3v4_nc", + "turboquant_3bit_nc", "int8_per_token_head", "fp8_per_token_head", ] diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 03a460fbe95..7028b12dab3 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1642,6 +1642,31 @@ class EngineArgs: 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) + 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, + ) + ray_runtime_env = None if is_ray_initialized(): # Ray Serve LLM calls `create_engine_config` in the context @@ -1948,6 +1973,19 @@ class EngineArgs: self.attention_backend ) + # TurboQuant requires FlashAttention 2 — FA3 boundary layers assert + # FlashAttentionImpl which fails with TurboQuantAttentionImpl. + if resolved_cache_dtype.startswith("turboquant_") and ( + attention_config.flash_attn_version is None + or attention_config.flash_attn_version >= 3 + ): + logger.warning( + "TurboQuant is not yet compatible with FlashAttention >= 3. " + "Overriding flash_attn_version to 2. To silence this " + "warning, pass --attention-config.flash_attn_version=2" + ) + attention_config.flash_attn_version = 2 + # Mamba config overrides mamba_config = copy.deepcopy(self.mamba_config) # Convert string to enum if needed (CLI parsing returns a string) diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 5e8825e2baf..a92e2f4ad18 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -379,6 +379,10 @@ class Attention(nn.Module, AttentionLayerBase): # Initialize KV cache quantization attributes _init_kv_cache_quant(self, quant_config, prefix) + # Initialize TurboQuant buffers (Pi, S, centroids) if tq cache dtype + if kv_cache_dtype.startswith("turboquant_"): + self._init_turboquant_buffers(kv_cache_dtype, head_size, prefix) + # for attn backends supporting query quantization self.query_quant = None if ( @@ -397,6 +401,67 @@ class Attention(nn.Module, AttentionLayerBase): else GroupShape.PER_TENSOR, ) + def _init_turboquant_buffers( + self, cache_dtype: str, head_size: int, prefix: str + ) -> None: + """Initialize TurboQuant rotation/projection matrices and centroids.""" + from vllm.model_executor.layers.quantization.turboquant.centroids import ( + get_centroids, + ) + from vllm.model_executor.layers.quantization.turboquant.config import ( + TurboQuantConfig, + ) + from vllm.model_executor.layers.quantization.turboquant.quantizer import ( + generate_wht_signs, + ) + + tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype, head_size) + + # Each layer needs a unique rotation matrix so quantization errors + # don't correlate across layers. Stride must exceed max head_dim to + # ensure non-overlapping RNG streams between adjacent layers. + _TQ_LAYER_SEED_STRIDE = 1337 + + from vllm.model_executor.models.utils import extract_layer_index + + layer_idx = extract_layer_index(prefix) + seed = tq_config.seed + layer_idx * _TQ_LAYER_SEED_STRIDE + + self.register_buffer( + "_tq_signs", + generate_wht_signs(head_size, seed=seed), + ) + self.register_buffer( + "_tq_centroids", + get_centroids(head_size, tq_config.centroid_bits), + ) + self._tq_config = tq_config + + # Pre-allocate decode intermediate buffers so model.to(device) moves + # them to GPU *before* the memory profiler runs. Without this the + # profiler gives all free memory to KV cache blocks and the first + # decode OOMs when these buffers are lazily allocated. + _vllm_cfg = get_current_vllm_config() + B = _vllm_cfg.scheduler_config.max_num_seqs + Hq = self.num_heads + S = _vllm_cfg.attention_config.tq_max_kv_splits_for_cuda_graph + D = head_size + self.register_buffer( + "_tq_mid_o_buf", + torch.empty(B, Hq, S, D + 1, dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "_tq_output_buf", + torch.empty(B, Hq, D, dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "_tq_lse_buf", + torch.empty(B, Hq, dtype=torch.float32), + persistent=False, + ) + def forward( self, query: torch.Tensor, @@ -544,6 +609,23 @@ class Attention(nn.Module, AttentionLayerBase): kv_quant_mode=quant_mode, sliding_window=self.sliding_window, ) + elif self.kv_cache_dtype.startswith("turboquant_"): + from vllm.model_executor.layers.quantization.turboquant.config import ( + TurboQuantConfig, + ) + from vllm.v1.kv_cache_interface import TQFullAttentionSpec + + tq_config = TurboQuantConfig.from_cache_dtype( + self.kv_cache_dtype, self.head_size + ) + return TQFullAttentionSpec( + block_size=block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + head_size_v=self.head_size, + dtype=self.kv_cache_torch_dtype, + tq_slot_size=tq_config.slot_size_aligned, + ) else: return FullAttentionSpec( block_size=block_size, diff --git a/vllm/model_executor/layers/quantization/turboquant/__init__.py b/vllm/model_executor/layers/quantization/turboquant/__init__.py new file mode 100644 index 00000000000..10ee032c9ec --- /dev/null +++ b/vllm/model_executor/layers/quantization/turboquant/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TurboQuant: Near-optimal KV-cache quantization for vLLM. + +PolarQuant compression: random rotation + per-coordinate Lloyd-Max +scalar quantization for keys, uniform quantization for values. + +Reference: "TurboQuant: Online Vector Quantization with Near-optimal +Distortion Rate" (ICLR 2026), Zandieh et al. +""" + +from vllm.model_executor.layers.quantization.turboquant.config import TurboQuantConfig + +__all__ = ["TurboQuantConfig"] diff --git a/vllm/model_executor/layers/quantization/turboquant/centroids.py b/vllm/model_executor/layers/quantization/turboquant/centroids.py new file mode 100644 index 00000000000..490265747c5 --- /dev/null +++ b/vllm/model_executor/layers/quantization/turboquant/centroids.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lloyd-Max optimal scalar quantizer for TurboQuant. + +After rotating a d-dimensional unit vector by a random orthogonal matrix, +each coordinate approximately follows N(0, 1/d) for d >= 64. +We solve the Lloyd-Max conditions to find optimal centroids. + +Based on: turboquant-pytorch/lloyd_max.py (Zandieh et al.) +""" + +import math +from functools import lru_cache + +import torch + + +def _gaussian_pdf(x: float, sigma2: float) -> float: + return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp(-x * x / (2 * sigma2)) + + +def _trapz(f, a: float, b: float, n: int = 200) -> float: + """Trapezoidal numerical integration (replaces scipy.integrate.quad).""" + h = (b - a) / n + result = 0.5 * (f(a) + f(b)) + for i in range(1, n): + result += f(a + i * h) + return result * h + + +def solve_lloyd_max( + d: int, + bits: int, + max_iter: int = 200, + tol: float = 1e-10, +) -> tuple[torch.Tensor, torch.Tensor]: + """Solve Lloyd-Max optimal quantizer for N(0, 1/d) distribution. + + Args: + d: Vector dimension (determines variance = 1/d). + bits: Number of quantization bits. + max_iter: Maximum Lloyd-Max iterations. + tol: Convergence tolerance. + + Returns: + centroids: Sorted tensor of 2^bits optimal centroids. + boundaries: Sorted tensor of 2^bits - 1 decision boundaries. + """ + n_levels = 2**bits + sigma2 = 1.0 / d + sigma = math.sqrt(sigma2) + + def pdf(x): + return _gaussian_pdf(x, sigma2) + + lo, hi = -3.5 * sigma, 3.5 * sigma + centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)] + + for _ in range(max_iter): + boundaries = [ + (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1) + ] + edges = [lo * 3] + boundaries + [hi * 3] + new_centroids = [] + for i in range(n_levels): + a, b = edges[i], edges[i + 1] + num = _trapz(lambda x: x * pdf(x), a, b) + den = _trapz(pdf, a, b) + new_centroids.append(num / den if den > 1e-15 else centroids[i]) + + if max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels)) < tol: + break + centroids = new_centroids + + boundaries = [(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)] + return ( + torch.tensor(centroids, dtype=torch.float32), + torch.tensor(boundaries, dtype=torch.float32), + ) + + +@lru_cache(maxsize=32) +def get_centroids(d: int, bits: int) -> torch.Tensor: + """Get precomputed Lloyd-Max centroids (cached).""" + centroids, _ = solve_lloyd_max(d, bits) + return centroids diff --git a/vllm/model_executor/layers/quantization/turboquant/config.py b/vllm/model_executor/layers/quantization/turboquant/config.py new file mode 100644 index 00000000000..289bed12077 --- /dev/null +++ b/vllm/model_executor/layers/quantization/turboquant/config.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TurboQuant configuration.""" + +import math +from dataclasses import dataclass + +# Named TQ presets: each maps to frozen config parameters. +# key_quant_bits: 8 = FP8 keys, 3-4 = MSE (Lloyd-Max) quantized keys. +# value_quant_bits: 3-4 = uniform quantized values. +TQ_PRESETS: dict[str, dict] = { + "turboquant_k8v4": { + "key_quant_bits": 8, + "value_quant_bits": 4, + "norm_correction": False, + }, + "turboquant_4bit_nc": { + "key_quant_bits": 4, + "value_quant_bits": 4, + "norm_correction": True, + }, + "turboquant_k3v4_nc": { + "key_quant_bits": 3, + "value_quant_bits": 4, + "norm_correction": True, + }, + "turboquant_3bit_nc": { + "key_quant_bits": 3, + "value_quant_bits": 3, + "norm_correction": True, + }, +} + + +@dataclass +class TurboQuantConfig: + """Configuration for TurboQuant KV-cache quantization. + + Uses PolarQuant (WHT rotation + Lloyd-Max scalar quantization) for keys + and uniform quantization for values. QJL is intentionally omitted — + community consensus (5+ independent groups) found it hurts attention + quality by amplifying variance through softmax. + + Named presets (use via --kv-cache-dtype): + turboquant_k8v4: FP8 keys + 4-bit values, 2.6x, +1.17% PPL + turboquant_4bit_nc: 4-bit MSE keys + 4-bit values + NC, 3.8x, +2.71% + turboquant_k3v4_nc: 3-bit MSE keys + 4-bit values + NC, ~3.5x, +10.63% + turboquant_3bit_nc: 3-bit MSE keys + 3-bit values + NC, 4.9x, +20.59% + + Args: + head_dim: Attention head dimension (e.g. 64, 96, 128). + key_quant_bits: Bits for key quantization. 8 = FP8 keys (no + rotation/MSE). 3-4 = Lloyd-Max MSE quantized keys. + value_quant_bits: Bits per value dimension for uniform quantization. + 3 = 8 levels, 4 = 16 levels (default). + seed: Base seed for deterministic random matrix generation. + Actual seed per layer = seed + layer_idx * 1337. + norm_correction: Re-normalize centroid vectors to unit norm before + inverse rotation during dequant. Fixes quantization-induced norm + distortion, improving PPL by ~0.8% at 4-bit. + """ + + head_dim: int = 128 + key_quant_bits: int = 3 # 3-4 = MSE keys, 8 = FP8 keys + value_quant_bits: int = 4 # 3-4 = uniform quantized values + seed: int = 42 + norm_correction: bool = False + + @property + def key_fp8(self) -> bool: + """Whether keys are stored as FP8 — no rotation/quantization needed.""" + return self.key_quant_bits == 8 + + @property + def mse_bits(self) -> int: + """MSE quantizer bit-width (determines centroid count: 2^mse_bits). + + For MSE key modes, equals key_quant_bits. + For FP8 key mode, falls back to value_quant_bits (centroids are still + needed for continuation-prefill dequant and decode kernel params). + """ + if self.key_fp8: + return self.value_quant_bits + return self.key_quant_bits + + @property + def key_mse_bits(self) -> int: + """MSE bits actually used for key quantization (0 if FP8 keys).""" + if self.key_fp8: + return 0 + return self.key_quant_bits + + @property + def centroid_bits(self) -> int: + """Bits for centroid generation — always non-zero.""" + return self.mse_bits + + @property + def n_centroids(self) -> int: + return 2**self.mse_bits + + @property + def key_packed_size(self) -> int: + """Packed bytes for a single KEY vector. + + FP8 mode (key_quant_bits=8): + head_dim bytes (1 byte per element, no overhead). + + TQ mode: + - MSE indices: ceil(head_dim * key_mse_bits / 8) bytes + - vec_norm: 2 bytes (float16) + """ + if self.key_fp8: + return self.head_dim # 1 byte per element + mse_bytes = math.ceil(self.head_dim * self.key_mse_bits / 8) + norm_bytes = 2 # vec_norm fp16 + return mse_bytes + norm_bytes + + @property + def effective_value_quant_bits(self) -> int: + """Actual bits used for value storage.""" + return self.value_quant_bits + + @property + def value_packed_size(self) -> int: + """Packed bytes for a single VALUE vector. + + Uniform quantization: ceil(head_dim * bits / 8) + 4 bytes (scale + zero fp16). + """ + data_bytes = math.ceil(self.head_dim * self.value_quant_bits / 8) + return data_bytes + 4 # +2 scale(fp16) +2 zero(fp16) + + @property + def slot_size(self) -> int: + """Total packed bytes per head per position (key + value combined). + + Layout: [key_packed | value_packed] + """ + return self.key_packed_size + self.value_packed_size + + @property + def slot_size_aligned(self) -> int: + """Slot size rounded up to next even number. + + Even-number is required so effective_head_size = slot_size_aligned // 2 + is integral. + """ + s = self.slot_size + 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. + """ + if n <= 0 or num_layers <= 0: + return [] + n = min(n, num_layers // 2) # don't skip more than half + first = list(range(n)) + last = list(range(num_layers - n, num_layers)) + # Deduplicate (if num_layers <= 2*n) + indices = sorted(set(first + last)) + return [str(i) for i in indices] + + @staticmethod + 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. + """ + if cache_dtype not in TQ_PRESETS: + valid = ", ".join(TQ_PRESETS.keys()) + raise ValueError( + f"Unknown TurboQuant cache dtype: {cache_dtype!r}. " + f"Valid presets: {valid}" + ) + preset = TQ_PRESETS[cache_dtype] + return TurboQuantConfig( + head_dim=head_dim, + key_quant_bits=preset["key_quant_bits"], + value_quant_bits=preset["value_quant_bits"], + norm_correction=preset["norm_correction"], + ) diff --git a/vllm/model_executor/layers/quantization/turboquant/quantizer.py b/vllm/model_executor/layers/quantization/turboquant/quantizer.py new file mode 100644 index 00000000000..aea63c52bac --- /dev/null +++ b/vllm/model_executor/layers/quantization/turboquant/quantizer.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TurboQuant quantizer utilities. + +Serving path uses generate_wht_signs() for WHT rotation sign buffers. +Triton kernels handle all quantization, packing, and dequantization on GPU. +""" + +import torch + +_CPU = torch.device("cpu") + + +def generate_wht_signs(d: int, seed: int, device: torch.device = _CPU) -> torch.Tensor: + """Generate deterministic random ±1 signs for WHT rotation. + + Used with Walsh-Hadamard Transform for per-layer rotation randomization. + Same seed derivation as QR (per-layer via seed + layer_idx * stride). + """ + gen = torch.Generator(device="cpu") + gen.manual_seed(seed) + bits = torch.randint(0, 2, (d,), generator=gen, device="cpu") + signs = bits.float() * 2 - 1 + return signs.to(device) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 045298dbb36..bbbd9af0cf7 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -255,6 +255,11 @@ class CudaPlatformBase(Platform): valid_backends_priorities = [] invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {} + # TurboQuant KV cache: route directly to TQ backend + kv_cache_dtype = attn_selector_config.kv_cache_dtype + if kv_cache_dtype is not None and kv_cache_dtype.startswith("turboquant_"): + return [(AttentionBackendEnum.TURBOQUANT, 0)], {} + backend_priorities = _get_backend_priorities( attn_selector_config.use_mla, device_capability, diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 892fc0c950b..f9a4f8bd732 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -61,6 +61,12 @@ class XPUPlatform(Platform): "only NHD layout is supported by XPU attention kernels." ) + # TurboQuant KV cache: route directly to TQ backend + kv_cache_dtype = attn_selector_config.kv_cache_dtype + if kv_cache_dtype is not None and kv_cache_dtype.startswith("turboquant_"): + logger.info_once("Using TurboQuant attention backend.") + return AttentionBackendEnum.TURBOQUANT.get_path() + dtype = attn_selector_config.dtype if attn_selector_config.use_sparse: logger.info_once("Using XPU MLA Sparse backend.") diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 60b40855b5d..26e377de69c 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -42,6 +42,10 @@ STR_DTYPE_TO_TORCH_DTYPE = { "fp8_per_token_head": torch.uint8, "fp8_inc": torch.float8_e4m3fn, "fp8_ds_mla": torch.uint8, + "turboquant_k8v4": torch.uint8, + "turboquant_4bit_nc": torch.uint8, + "turboquant_k3v4_nc": torch.uint8, + "turboquant_3bit_nc": torch.uint8, } TORCH_DTYPE_TO_NUMPY_DTYPE = { diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 4744ead4f54..f31edfafc38 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -82,6 +82,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "RocmAiterUnifiedAttentionBackend" ) CPU_ATTN = "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend" + TURBOQUANT = "vllm.v1.attention.backends.turboquant_attn.TurboQuantAttentionBackend" # Placeholder for third-party/custom backends - must be registered before use # set to None to avoid alias with other backend, whose value is an empty string CUSTOM = None diff --git a/vllm/v1/attention/backends/turboquant_attn.py b/vllm/v1/attention/backends/turboquant_attn.py new file mode 100644 index 00000000000..279fcb04ace --- /dev/null +++ b/vllm/v1/attention/backends/turboquant_attn.py @@ -0,0 +1,812 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TurboQuant attention backend for vLLM. + +Prefill: Standard scaled dot-product attention on uncompressed K/V, + then quantize K and store K+V into combined cache slot. +Decode: Compute TQ attention scores from compressed cache, + unpack FP16 values, softmax + weighted sum. + +Cache layout (no leading 2 dimension): + (num_blocks, block_size, num_kv_heads, slot_size) + where slot_size = key_packed_size + value_fp16_size + +Per-head per-position slot layout: + [key_packed (kps bytes) | value_fp16 (D*2 bytes)] + For turboquant_k3v4_nc head_dim=256: [100 bytes key | 512 bytes value] = 612 +""" + +import functools +import math +from dataclasses import dataclass +from typing import Any, ClassVar + +import torch +import torch.nn.functional as F + +from vllm.config import get_current_vllm_config +from vllm.config.cache import CacheDType +from vllm.triton_utils import triton +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionImpl, + AttentionLayer, + AttentionMetadata, + AttentionMetadataBuilder, + AttentionType, + CommonAttentionMetadata, + MultipleOf, +) +from vllm.v1.attention.backends.fa_utils import ( + is_flash_attn_varlen_func_available, +) +from vllm.v1.attention.backends.utils import split_decodes_and_prefills +from vllm.v1.attention.ops.triton_turboquant_decode import ( + _tq_full_dequant_kv, + _use_fp8_e4b15, + triton_turboquant_decode_attention, +) +from vllm.v1.attention.ops.triton_turboquant_store import triton_turboquant_store + +_HAS_FLASH_ATTN = is_flash_attn_varlen_func_available() +if _HAS_FLASH_ATTN: + from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func + +# Continuation prefill: for small continuation chunks (q_len ≤ threshold), +# use the TQ decode kernel directly instead of full-dequant + flash_attn. +# do_kv_cache_update already stored all tokens to TQ cache, so the decode +# kernel can read them efficiently. This avoids O(cached_len) dequant work +# per continuation, eliminating the O(N²/chunk_size) collapse at long context. +_CONTINUATION_DECODE_THRESHOLD = 128 + + +def _build_hadamard(d: int, device_str: str) -> torch.Tensor: + """Orthonormal Hadamard matrix (Sylvester construction), cached per (d, device). + + Precomputed D×D matrix enables matmul-based WHT — single cuBLAS GEMM + instead of log2(D) butterfly kernel launches. 64KB for D=128. + """ + # Normalize device string so "cuda" and "cuda:0" hit the same cache entry. + return _build_hadamard_cached(d, str(torch.device(device_str))) + + +@functools.cache +def _build_hadamard_cached(d: int, device_str: str) -> torch.Tensor: + H = torch.tensor([[1.0]]) + while H.shape[0] < d: + H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0) + return (H / math.sqrt(d)).to(torch.device(device_str)) + + +class TurboQuantAttentionBackend(AttentionBackend): + """Attention backend using TurboQuant KV-cache compression.""" + + accept_output_buffer: bool = True + forward_includes_kv_cache_update: bool = False + + supported_dtypes: ClassVar[list[torch.dtype]] = [ + torch.float16, + torch.bfloat16, + ] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "turboquant_k8v4", + "turboquant_4bit_nc", + "turboquant_k3v4_nc", + "turboquant_3bit_nc", + ] + + @staticmethod + def get_name() -> str: + return "TURBOQUANT" + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [16, 32, 64, 128] + + @classmethod + def supports_attn_type(cls, attn_type: str) -> bool: + return attn_type == AttentionType.DECODER + + @classmethod + def supports_per_head_quant_scales(cls) -> bool: + return False + + @staticmethod + def get_impl_cls() -> type["TurboQuantAttentionImpl"]: + return TurboQuantAttentionImpl + + @staticmethod + def get_builder_cls() -> type["TurboQuantMetadataBuilder"]: + return TurboQuantMetadataBuilder + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "turboquant_4bit_nc", + ) -> tuple[int, ...]: + """Combined K+V cache shape — no leading 2 dimension. + + Standard attention backends use (2, num_blocks, block_size, num_kv_heads, + head_dim) with a leading 2 to separate K and V. TurboQuant packs K+V + into a single interleaved slot per head per position, so the cache is: + + (num_blocks, block_size, num_kv_heads, slot_size_aligned) + + Each slot = [key_packed | value_packed | padding]. + This is safe because TQ has its own get_kv_cache_shape override and + never shares cache tensors with other backends. Layers that fall back + to native dtype via kv_cache_dtype_skip_layers get their own + standard-shaped cache allocation. + + head_size is the model's real head_dim. slot_size_aligned is computed + from the TQ config to ensure correct cache allocation for all head dims. + """ + from vllm.model_executor.layers.quantization.turboquant.config import ( + TurboQuantConfig, + ) + + tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype_str, head_size) + return (num_blocks, block_size, num_kv_heads, tq_config.slot_size_aligned) + + @classmethod + def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: + if kv_cache_dtype is None: + return False + return kv_cache_dtype.startswith("turboquant_") + + @classmethod + def supports_head_size(cls, head_size: int) -> bool: + # head_size from spec is effective_head_size (padded_slot//2), + # not the model's actual head_dim. Accept any positive value. + return head_size > 0 + + +@dataclass +class TurboQuantMetadata(AttentionMetadata): + """Metadata for TurboQuant attention.""" + + seq_lens: torch.Tensor # (num_reqs,) — total context length per request + slot_mapping: torch.Tensor # (num_tokens,) — cache slot for each token + block_table: torch.Tensor # (num_reqs, max_num_blocks) + query_start_loc: torch.Tensor # (num_reqs + 1,) — cu_seqlens for queries + num_actual_tokens: int = 0 # actual tokens (excluding padding) + max_query_len: int = 0 # longest query in batch + max_seq_len: int = 0 # longest context in batch + is_prefill: bool = False + num_decodes: int = 0 # number of decode requests (first in batch) + num_decode_tokens: int = 0 # tokens from decode requests + + +class TurboQuantMetadataBuilder(AttentionMetadataBuilder[TurboQuantMetadata]): + """Builds TurboQuantMetadata from scheduler output.""" + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + def __init__(self, kv_cache_spec, layer_names, vllm_config, device): + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self._init_reorder_batch_threshold(1, supports_spec_as_decode=False) + + def build_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata + ) -> TurboQuantMetadata: + attn_metadata = self.build(0, common_attn_metadata) + # Set seq_lens to 1 so CUDA graph capture is fast + # (real seq_lens are filled at replay time). + attn_metadata.seq_lens.fill_(1) + return attn_metadata + + def build(self, common_prefix_len, common_attn_metadata, fast_build=False): + """Build TurboQuantMetadata from common attention metadata.""" + cam = common_attn_metadata + + # With reorder_batch_threshold=1, the model runner guarantees + # decodes come first in the batch. split_decodes_and_prefills + # finds the boundary (operates on CPU tensors — no GPU sync). + assert self.reorder_batch_threshold is not None + num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills( + cam, decode_threshold=self.reorder_batch_threshold + ) + + return TurboQuantMetadata( + seq_lens=cam.seq_lens, + slot_mapping=cam.slot_mapping, + block_table=cam.block_table_tensor, + query_start_loc=cam.query_start_loc, + num_actual_tokens=cam.num_actual_tokens, + max_query_len=cam.max_query_len, + max_seq_len=cam.max_seq_len, + is_prefill=(cam.max_query_len > 1), + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + ) + + +class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]): + """TurboQuant attention implementation. + + Vectorized PyTorch: batch quantize/store, vectorized bit-unpack + decode with einsum scores and value gather. + """ + + supports_quant_query_input: bool = False + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int | None = None, + alibi_slopes: list[float] | None = None, + sliding_window: int | None = None, + kv_cache_dtype: str = "auto", + logits_soft_cap: float | None = None, + attn_type: str = AttentionType.DECODER, + kv_sharing_target_layer_name: str | None = None, + **kwargs, + ): + self.num_heads = num_heads + self.head_size = head_size + self.scale = scale + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.kv_cache_dtype = kv_cache_dtype + + from vllm.model_executor.layers.quantization.turboquant.config import ( + TurboQuantConfig, + ) + + self.tq_config = TurboQuantConfig.from_cache_dtype(kv_cache_dtype, head_size) + + # Pre-compute kernel constants from config (avoid repeated arithmetic) + cfg = self.tq_config + self._mse_bytes = ( + math.ceil(head_size * cfg.key_mse_bits / 8) + if not cfg.key_fp8 + else head_size + ) + self._val_data_bytes = math.ceil(head_size * cfg.effective_value_quant_bits / 8) + self._n_centroids = cfg.n_centroids if not cfg.key_fp8 else 1 + + # Fixed NUM_KV_SPLITS (grid dims must be constant for cudagraph, + # and benchmarks show no regression vs dynamic in eager mode). + vllm_config = get_current_vllm_config() + self.max_num_kv_splits = ( + vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph + ) + + def _ensure_on_device(self, layer, device): + """One-time derivation of TQ buffers (rotation matrices, midpoints). + + Registered buffers (_tq_signs, _tq_centroids) are already on the + correct device via register_buffer + model.to(device). + """ + if not hasattr(layer, "_tq_cached"): + D = layer._tq_signs.shape[0] + signs = layer._tq_signs.to(device=device, dtype=torch.float32) + + # WHT rotation: orthonormal + self-inverse, enabling future + # in-kernel butterfly fusion and trivial inverse for continuation. + H = _build_hadamard(D, str(device)) + layer._tq_PiT = (signs.unsqueeze(1) * H).contiguous() + layer._tq_Pi = layer._tq_PiT.T.contiguous() + + c = layer._tq_centroids.to(device=device, dtype=torch.float32) + # Precompute midpoints for threshold-based quantization + c_sorted, _ = c.sort() + layer._tq_midpoints = (c_sorted[:-1] + c_sorted[1:]) / 2 + # Decode buffers (_tq_mid_o_buf, _tq_output_buf, _tq_lse_buf) + # are pre-allocated via register_buffer in Attention.__init__ + # and moved to GPU by model.to(device) — no allocation needed + # here. The memory profiler sees them before KV cache sizing. + layer._tq_cached = True + + def do_kv_cache_update( + self, + layer: torch.nn.Module, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> None: + """Store compressed K/V into the combined TQ cache. + + Called as a separate custom op (unified_kv_cache_update) BEFORE + the attention forward, matching FlashAttention's split pattern. + slot_mapping is already sliced to num_actual_tokens by the caller. + """ + N = slot_mapping.shape[0] + if N <= 0: + return + + device = key.device + self._ensure_on_device(layer, device) + + k = key[:N].view(N, self.num_kv_heads, self.head_size) + v = value[:N].view(N, self.num_kv_heads, self.head_size) + self._store_kv(k, v, kv_cache, slot_mapping, layer) + + def forward( + self, + layer: AttentionLayer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: "TurboQuantMetadata", + output: torch.Tensor | None = None, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + num_tokens = query.shape[0] + + if output is None: + output = torch.zeros( + num_tokens, + self.num_heads * self.head_size, + dtype=query.dtype, + device=query.device, + ) + + if attn_metadata is None: + return output.fill_(0) + + # Slice to actual tokens + N = attn_metadata.num_actual_tokens + if N <= 0: + return output.fill_(0) + + q = query[:N].view(N, self.num_heads, self.head_size) + + # Get TQ buffers, ensure on device (one-time migration). + # Use Any-typed alias for dynamic _tq_* attrs set by _ensure_on_device. + tq_layer: Any = layer + device = q.device + self._ensure_on_device(tq_layer, device) + Pi = tq_layer._tq_Pi + PiT = tq_layer._tq_PiT + centroids = tq_layer._tq_centroids + + # Compute attention (KV cache was already updated by do_kv_cache_update) + # With reorder_batch_threshold=1, decodes come first in the batch. + # num_decodes/num_decode_tokens from metadata give the split point. + num_decodes = attn_metadata.num_decodes + num_decode_tokens = attn_metadata.num_decode_tokens + + if not attn_metadata.is_prefill: + # Pure decode batch — fast path + attn_out = self._decode_attention( + q, kv_cache, attn_metadata, Pi, centroids, PiT, layer + ) + elif num_decodes == 0: + # Pure prefill batch + k = key[:N].view(N, self.num_kv_heads, self.head_size) + v = value[:N].view(N, self.num_kv_heads, self.head_size) + attn_out = self._prefill_attention( + q, + k, + v, + kv_cache, + attn_metadata, + Pi, + centroids, + PiT, + layer=layer, + ) + else: + # Mixed batch: decodes first (guaranteed by reorder_batch). + attn_out = torch.zeros( + N, self.num_heads, self.head_size, device=device, dtype=q.dtype + ) + + # --- Decode portion (first num_decodes requests) --- + # Use full-batch max_seq_len as safe upper bound (no GPU sync). + decode_meta = TurboQuantMetadata( + seq_lens=attn_metadata.seq_lens[:num_decodes], + slot_mapping=attn_metadata.slot_mapping[:num_decode_tokens], + block_table=attn_metadata.block_table[:num_decodes], + query_start_loc=attn_metadata.query_start_loc[: num_decodes + 1], + num_actual_tokens=num_decode_tokens, + max_query_len=1, + max_seq_len=attn_metadata.max_seq_len, + is_prefill=False, + ) + attn_out[:num_decode_tokens] = self._decode_attention( + q[:num_decode_tokens], kv_cache, decode_meta, Pi, centroids, PiT, layer + ) + + # --- Prefill portion (remaining requests) --- + # CRITICAL: use prefill-specific max_seq_len so flash_attn's + # fast path (max_query_len == max_seq_len) triggers for + # first-chunk prefills. Using full-batch max_seq_len breaks + # this because decode requests inflate max_seq_len. + prefill_seq_lens = attn_metadata.seq_lens[num_decodes:] + # Use CPU-side max to avoid GPU→CPU sync from .item() + prefill_max_seq = max(attn_metadata.seq_lens[num_decodes:].tolist()) + prefill_qsl = ( + attn_metadata.query_start_loc[num_decodes:] - num_decode_tokens + ) + prefill_meta = TurboQuantMetadata( + seq_lens=prefill_seq_lens, + slot_mapping=attn_metadata.slot_mapping[num_decode_tokens:N], + block_table=attn_metadata.block_table[num_decodes:], + query_start_loc=prefill_qsl, + num_actual_tokens=N - num_decode_tokens, + max_query_len=attn_metadata.max_query_len, + max_seq_len=prefill_max_seq, + is_prefill=True, + ) + k = key[:N].view(N, self.num_kv_heads, self.head_size) + v = value[:N].view(N, self.num_kv_heads, self.head_size) + attn_out[num_decode_tokens:] = self._prefill_attention( + q[num_decode_tokens:], + k[num_decode_tokens:], + v[num_decode_tokens:], + kv_cache, + prefill_meta, + Pi, + centroids, + PiT, + layer=layer, + ) + + # Write into output buffer: attn_out is (N, Hq, D) + # output may be 2D (N, Hq*D) or 3D (N, Hq, D) + if output.ndim == 3: + output[:N] = attn_out.to(output.dtype) + else: + output[:N] = attn_out.reshape(N, -1).to(output.dtype) + return output + + # ------------------------------------------------------------------ # + # Store K/V into combined cache (vectorized) # + # ------------------------------------------------------------------ # + def _store_kv( + self, + key: torch.Tensor, # (N, Hk, D) + value: torch.Tensor, # (N, Hk, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + slot_mapping: torch.Tensor, + layer: Any, + ): + """Quantize + store via fused Triton kernel.""" + triton_turboquant_store( + key, + value, + kv_cache, + slot_mapping, + layer._tq_PiT, + layer._tq_midpoints, + mse_bits=self.tq_config.key_mse_bits, + key_packed_size=self.tq_config.key_packed_size, + value_quant_bits=self.tq_config.effective_value_quant_bits, + key_fp8=self.tq_config.key_fp8, + ) + + # ------------------------------------------------------------------ # + # Prefill: SDPA on raw Q/K/V with causal mask # + # ------------------------------------------------------------------ # + def _prefill_attention( + self, + query: torch.Tensor, # (N, Hq, D) + key: torch.Tensor, # (N, Hk, D) + value: torch.Tensor, # (N, Hk, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + attn_metadata: TurboQuantMetadata, + Pi: torch.Tensor, + centroids: torch.Tensor, + PiT: torch.Tensor | None = None, + layer: Any = None, + ) -> torch.Tensor: + N, Hq, D = query.shape + + # Fast path: use flash_attn for first-chunk prefills (all K/V in batch). + # max_query_len == max_seq_len means no request has prior cached KV. + # Both are Python ints — no GPU sync. + if _HAS_FLASH_ATTN and attn_metadata.max_query_len == attn_metadata.max_seq_len: + output = torch.empty(N, Hq, D, device=query.device, dtype=query.dtype) + flash_attn_varlen_func( + q=query, + k=key, + v=value, + cu_seqlens_q=attn_metadata.query_start_loc, + cu_seqlens_k=attn_metadata.query_start_loc, + max_seqlen_q=attn_metadata.max_query_len, + max_seqlen_k=attn_metadata.max_query_len, + softmax_scale=self.scale, + causal=True, + out=output, + ) + return output + + # Continuation or no flash_attn: per-request attention. + # For continuation chunks (seq_len > q_len), we must attend to + # previously cached K/V from the TQ cache, not just the current + # chunk's raw K/V. + Hk = key.shape[1] + use_gqa = Hk < Hq + query_start_loc = attn_metadata.query_start_loc + num_reqs = query_start_loc.shape[0] - 1 + + output = torch.zeros(N, Hq, D, device=query.device, dtype=query.dtype) + + # Convert to Python lists once (single CPU-GPU sync) instead of + # per-request .item() calls that each force a sync. + qsl = query_start_loc.tolist() + seq_lens_list = attn_metadata.seq_lens.tolist() + + # Pre-allocate cu_seqlens for single-request flash_attn calls + # to avoid per-request host→device tensor creation. + _cu_2 = torch.zeros(2, device=query.device, dtype=torch.int32) + + for i in range(num_reqs): + q_start = qsl[i] + q_end = qsl[i + 1] + q_len = q_end - q_start + if q_len <= 0: + continue + + seq_len = seq_lens_list[i] + q_seq = query[q_start:q_end] # (q_len, Hq, D) + k_seq = key[q_start:q_end] # (q_len, Hk, D) + v_seq = value[q_start:q_end] # (q_len, Hk, D) + + if q_len == seq_len: + # First-chunk prefill: all K/V are in the current batch. + if _HAS_FLASH_ATTN: + out = torch.empty_like(q_seq) + _cu_2[1] = q_len + cu = _cu_2 + flash_attn_varlen_func( + q=q_seq, + k=k_seq, + v=v_seq, + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=q_len, + max_seqlen_k=q_len, + softmax_scale=self.scale, + causal=True, + out=out, + ) + else: + q_t = q_seq.transpose(0, 1).contiguous() + k_t = k_seq.transpose(0, 1).contiguous() + v_t = v_seq.transpose(0, 1).contiguous() + out = F.scaled_dot_product_attention( + q_t, + k_t, + v_t, + is_causal=True, + scale=self.scale, + enable_gqa=use_gqa, + ).transpose(0, 1) + output[q_start:q_end] = out.to(query.dtype) + else: + # Continuation chunk: tokens already stored to TQ cache + # by do_kv_cache_update. Use decode kernel directly to + # avoid O(cached_len) full-dequant per continuation. + # For large continuations, fall back to _continuation_prefill. + cached_len = seq_len - q_len + if q_len <= _CONTINUATION_DECODE_THRESHOLD: + # Fast path: treat each query as a decode request + # with incremental seq_lens for causal masking. + synth_seq_lens = torch.arange( + cached_len + 1, + seq_len + 1, + device=query.device, + dtype=attn_metadata.seq_lens.dtype, + ) + synth_bt = attn_metadata.block_table[i : i + 1].expand(q_len, -1) + out = triton_turboquant_decode_attention( + query=q_seq, + kv_cache=kv_cache, + block_table=synth_bt, + seq_lens=synth_seq_lens, + Pi=Pi, + centroids=centroids, + scale=self.scale, + mse_bits=self.tq_config.key_mse_bits, + key_packed_size=self.tq_config.key_packed_size, + value_quant_bits=(self.tq_config.effective_value_quant_bits), + key_fp8=self.tq_config.key_fp8, + norm_correction=self.tq_config.norm_correction, + PiT=PiT, + ) + else: + # Large continuation: dequant cached K/V and use + # flash_attn for better throughput. + out = self._continuation_prefill( + layer, + q_seq, + k_seq, + v_seq, + kv_cache, + attn_metadata.block_table[i : i + 1], + cached_len, + seq_len, + Pi, + centroids, + ) + output[q_start:q_end] = out.to(query.dtype) + + return output + + def _continuation_prefill( + self, + layer: Any, + query: torch.Tensor, # (q_len, Hq, D) + key_chunk: torch.Tensor, # (q_len, Hk, D) + val_chunk: torch.Tensor, # (q_len, Hk, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + block_table: torch.Tensor, # (1, max_num_blocks) + cached_len: int, + seq_len: int, + Pi: torch.Tensor, + centroids: torch.Tensor, + ) -> torch.Tensor: + """Handle continuation chunk by dequanting cached K/V from TQ cache. + + Dequants previously cached K/V, concatenates with the current + chunk's raw K/V, then runs flash_attn with causal masking. + """ + q_len, Hq, D = query.shape + Hk = key_chunk.shape[1] + device = query.device + block_size = kv_cache.shape[1] + BLOCK_D = triton.next_power_of_2(D) + + mse_bytes = self._mse_bytes + val_data_bytes = self._val_data_bytes + + # Dequant cached K/V from TQ cache + # Allocate slightly over to align to block_size for the grid. + # Reuse cached buffers to avoid per-call allocation (~16MB at 8K). + alloc_len = math.ceil(cached_len / block_size) * block_size + buf_shape = (1, Hk, alloc_len, D) + k_buf = getattr(layer, "_tq_k_dequant_buf", None) + if k_buf is None or k_buf.shape[2] < alloc_len: + k_buf = torch.empty(buf_shape, dtype=torch.float16, device=device) + v_buf = torch.empty(buf_shape, dtype=torch.float16, device=device) + layer._tq_k_dequant_buf = k_buf + layer._tq_v_dequant_buf = v_buf + else: + v_buf = layer._tq_v_dequant_buf + k_cached = k_buf[:, :, :alloc_len, :].zero_() + v_cached = v_buf[:, :, :alloc_len, :].zero_() + + grid = (alloc_len, 1 * Hk) + _tq_full_dequant_kv[grid]( + kv_cache, + block_table, + centroids, + k_cached, + v_cached, + k_cached.stride(0), + k_cached.stride(1), + k_cached.stride(2), + v_cached.stride(0), + v_cached.stride(1), + v_cached.stride(2), + kv_cache.stride(0), + kv_cache.stride(1), + kv_cache.stride(2), + block_table.stride(0), + HEAD_DIM=D, + BLOCK_SIZE=block_size, + NUM_KV_HEADS=Hk, + MSE_BYTES=mse_bytes, + KPS=self.tq_config.key_packed_size, + VQB=self.tq_config.effective_value_quant_bits, + VAL_DATA_BYTES=val_data_bytes, + MSE_BITS=self.tq_config.key_mse_bits, + KEY_FP8=1 if self.tq_config.key_fp8 else 0, + BLOCK_D=BLOCK_D, + NORM_CORRECTION=1 if self.tq_config.norm_correction else 0, + FP8_E4B15=_use_fp8_e4b15(device.index or 0), + num_warps=4, + ) + + # Inverse-rotate MSE keys back to original space + if not self.tq_config.key_fp8: + k_flat = k_cached[0, :, :cached_len, :].reshape(-1, D).float() + k_flat = k_flat @ Pi + k_cached_trim = ( + k_flat.to(torch.float16).reshape(Hk, cached_len, D).transpose(0, 1) + ) # (cached_len, Hk, D) + else: + k_cached_trim = ( + k_cached[0, :, :cached_len, :].transpose(0, 1).contiguous() + ) # (cached_len, Hk, D) + + v_cached_trim = ( + v_cached[0, :, :cached_len, :].transpose(0, 1).contiguous() + ) # (cached_len, Hk, D) + + # Concatenate cached + current chunk K/V (match query dtype) + qdtype = query.dtype + k_full = torch.cat([k_cached_trim.to(qdtype), key_chunk], dim=0) + v_full = torch.cat([v_cached_trim.to(qdtype), val_chunk], dim=0) + + # Attention: q_len queries attending to seq_len K/V with causal mask + if _HAS_FLASH_ATTN: + output = torch.empty(q_len, Hq, D, device=device, dtype=query.dtype) + cu_seqlens_q = torch.tensor([0, q_len], device=device, dtype=torch.int32) + cu_seqlens_k = torch.tensor([0, seq_len], device=device, dtype=torch.int32) + flash_attn_varlen_func( + q=query, + k=k_full, + v=v_full, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=q_len, + max_seqlen_k=seq_len, + softmax_scale=self.scale, + causal=True, + out=output, + ) + return output + else: + # SDPA fallback: expand KV for GQA, build causal mask + q_t = query.transpose(0, 1).unsqueeze(0) # (1, Hq, q_len, D) + k_t = k_full.transpose(0, 1).unsqueeze(0) # (1, Hk, seq_len, D) + v_t = v_full.transpose(0, 1).unsqueeze(0) # (1, Hk, seq_len, D) + # Build causal mask: query position p can attend to K position j + # where j <= cached_len + p (p is 0-indexed within chunk) + q_pos = torch.arange(q_len, device=device).unsqueeze(1) + cached_len + k_pos = torch.arange(seq_len, device=device).unsqueeze(0) + mask = k_pos <= q_pos # (q_len, seq_len) + out = F.scaled_dot_product_attention( + q_t, + k_t, + v_t, + attn_mask=mask, + scale=self.scale, + enable_gqa=(Hk < Hq), + ) # (1, Hq, q_len, D) + return out[0].transpose(0, 1) # (q_len, Hq, D) + + # ------------------------------------------------------------------ # + # Decode: Triton TQ decode attention # + # ------------------------------------------------------------------ # + def _decode_attention( + self, + query: torch.Tensor, # (B, Hq, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + attn_metadata: TurboQuantMetadata, + Pi: torch.Tensor, + centroids: torch.Tensor, + PiT: torch.Tensor | None = None, + layer: torch.nn.Module | None = None, + ) -> torch.Tensor: + # Grab cached decode buffers from the layer (lazily allocated). + mid_o_buf = output_buf = lse_buf = None + if layer is not None: + mid_o_buf = getattr(layer, "_tq_mid_o_buf", None) + output_buf = getattr(layer, "_tq_output_buf", None) + lse_buf = getattr(layer, "_tq_lse_buf", None) + + result = triton_turboquant_decode_attention( + query=query, + kv_cache=kv_cache, + block_table=attn_metadata.block_table, + seq_lens=attn_metadata.seq_lens, + Pi=Pi, + centroids=centroids, + scale=self.scale, + mse_bits=self.tq_config.key_mse_bits, + key_packed_size=self.tq_config.key_packed_size, + value_quant_bits=self.tq_config.effective_value_quant_bits, + key_fp8=self.tq_config.key_fp8, + norm_correction=self.tq_config.norm_correction, + PiT=PiT, + mid_o_buf=mid_o_buf, + output_buf=output_buf, + lse_buf=lse_buf, + buf_holder=layer, + max_num_kv_splits=self.max_num_kv_splits, + ) + return result diff --git a/vllm/v1/attention/ops/triton_turboquant_decode.py b/vllm/v1/attention/ops/triton_turboquant_decode.py new file mode 100644 index 00000000000..8b276e31eaf --- /dev/null +++ b/vllm/v1/attention/ops/triton_turboquant_decode.py @@ -0,0 +1,617 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton fused TurboQuant decode attention. + +Decode path: Triton stage1 (split-KV tiled attention scoring + value +accumulation) + stage2 (log-sum-exp reduction across splits). + +Supports FP8 (E4M3) keys, 3-bit and 4-bit uniform quantized values. +""" + +import math +from typing import Any + +import torch + +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.triton_decode_attention import ( + _fwd_kernel_stage2, +) + +_FP8_E4B15: dict[int, int] = {} + + +def _use_fp8_e4b15(device: int = 0) -> int: + """Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0.""" + if device not in _FP8_E4B15: + cap = torch.cuda.get_device_capability(device) + _FP8_E4B15[device] = 1 if cap < (8, 9) else 0 + return _FP8_E4B15[device] + + +# --------------------------------------------------------------------------- +# Stage 1: Fused TQ score + value accumulation (BLOCK_KV tiled) +# --------------------------------------------------------------------------- + + +@triton.jit +def _tq_decode_stage1( + # Precomputed query projection + Q_rot_ptr, # [B, Hq, D] float32 + # Compressed KV cache (combined K+V) + KV_cache_ptr, # [num_blocks, block_size, Hk, padded_slot] uint8 + # Block table and sequence info + Block_table_ptr, # [B, max_num_blocks] int32 + Seq_lens_ptr, # [B] int32 + # TQ parameters + Centroids_ptr, # [n_centroids] float32 + # Output (intermediate for stage2) + Mid_o_ptr, # [B, Hq, NUM_KV_SPLITS, D+1] float32 + # Strides + stride_qb, + stride_qh, # Q strides: [B, Hq, D] + stride_cache_block, + stride_cache_pos, + stride_cache_head, # KV cache + stride_bt_b, # block_table stride per batch + stride_mid_b, + stride_mid_h, + stride_mid_s, # mid_o strides + # Constexpr dims + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_SIZE: tl.constexpr, # KV cache block_size (pages) + NUM_KV_SPLITS: tl.constexpr, + KV_GROUP_SIZE: tl.constexpr, # Hq // Hk + # TQ layout constants + MSE_BITS: tl.constexpr, # 3 or 4 + MSE_BYTES: tl.constexpr, # ceil(D * mse_bits / 8) + KPS: tl.constexpr, # key_packed_size + VQB: tl.constexpr, # value_quant_bits (4 or 8=FP8) + VAL_DATA_BYTES: tl.constexpr, # ceil(D * vqb / 8) or D for FP8 + # Score constants + ATTN_SCALE: tl.constexpr, # 1/sqrt(D) + # Block tile sizes + BLOCK_D: tl.constexpr, # next_power_of_2(HEAD_DIM) + BLOCK_KV: tl.constexpr, # tokens per tile (16) + KEY_FP8: tl.constexpr, # 1 if K is stored as FP8 + NORM_CORRECTION: tl.constexpr = 0, # 1 = re-normalize centroids + FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+) +): + bid = tl.program_id(0) # batch index + hid = tl.program_id(1) # q_head index + sid = tl.program_id(2) # kv_split index + + kv_head = hid // KV_GROUP_SIZE + + # Sequence length for this batch + seq_len = tl.load(Seq_lens_ptr + bid) + + # KV split range + split_len = tl.cdiv(seq_len, NUM_KV_SPLITS) + split_start = split_len * sid + split_end = tl.minimum(split_start + split_len, seq_len) + + if split_start >= split_end: + return + + # Dimension offsets + d_offs = tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + kv_range = tl.arange(0, BLOCK_KV) + + # Load query vector: q_rot — [BLOCK_D] float32 + q_base = bid * stride_qb + hid * stride_qh + q_rot = tl.load(Q_rot_ptr + q_base + d_offs, mask=d_mask, other=0.0).to(tl.float32) + + # Precompute byte/bit index vectors for MSE gather loads + if not KEY_FP8: + mse_bit_off = d_offs * MSE_BITS + mse_byte_idx = mse_bit_off // 8 + mse_bit_shift = mse_bit_off % 8 + mse_mask = (1 << MSE_BITS) - 1 + + # Precompute value bit/byte index vectors (loop-invariant) + if VQB == 3: + val_bit_off = d_offs * 3 + val_byte_idx = val_bit_off // 8 + val_bit_shift = val_bit_off % 8 + + # Online softmax accumulators + m_prev = -float("inf") + l_prev = 0.0 + acc = tl.zeros([BLOCK_D], dtype=tl.float32) + + bt_base = bid * stride_bt_b + + # ================================================================ + # TILED LOOP: process BLOCK_KV tokens per iteration + # ================================================================ + for start_n in range(split_start, split_end, BLOCK_KV): + kv_offs = start_n + kv_range + kv_mask = kv_offs < split_end + + page_idx = kv_offs // BLOCK_SIZE + page_off = kv_offs % BLOCK_SIZE + block_nums = tl.load( + Block_table_ptr + bt_base + page_idx, + mask=kv_mask, + other=0, + ) + + slot_bases = ( + block_nums * stride_cache_block + + page_off * stride_cache_pos + + kv_head * stride_cache_head + ) + + # ============================================================ + # COMPUTE ATTENTION SCORES: [BLOCK_KV] + # ============================================================ + if KEY_FP8: + k_addrs = slot_bases[:, None] + d_offs[None, :] + k_raw = tl.load( + KV_cache_ptr + k_addrs, + mask=kv_mask[:, None] & d_mask[None, :], + other=0, + ) + if FP8_E4B15: + k_float = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32) + else: + k_float = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32) + scores = ( + tl.sum( + tl.where(d_mask[None, :], q_rot[None, :] * k_float, 0.0), + axis=1, + ) + * ATTN_SCALE + ) + scores = tl.where(kv_mask, scores, -float("inf")) + else: + # MSE unpack + norms + mse_addrs0 = slot_bases[:, None] + mse_byte_idx[None, :] + mse_raw0 = tl.load( + KV_cache_ptr + mse_addrs0, + mask=kv_mask[:, None] & d_mask[None, :], + other=0, + ).to(tl.int32) + mse_raw1 = tl.load( + KV_cache_ptr + mse_addrs0 + 1, + mask=kv_mask[:, None] & d_mask[None, :], + other=0, + ).to(tl.int32) + raw16 = mse_raw0 | (mse_raw1 << 8) + mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask + + # Centroid gather + dot product + c_vals = tl.load( + Centroids_ptr + mse_idx, + mask=kv_mask[:, None] & d_mask[None, :], + other=0.0, + ) + + # Norm correction: re-normalize centroid vector to unit norm + if NORM_CORRECTION: + c_norm_sq = tl.sum( + tl.where(d_mask[None, :], c_vals * c_vals, 0.0), + axis=1, + ) + c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16) + c_vals = c_vals * c_inv_norm[:, None] + + term1 = tl.sum( + tl.where(d_mask[None, :], q_rot[None, :] * c_vals, 0.0), + axis=1, + ) + + # Load norms (fp16 -> fp32): norms are at MSE_BYTES offset + norm_bases = slot_bases + MSE_BYTES + n_lo = tl.load(KV_cache_ptr + norm_bases, mask=kv_mask, other=0).to( + tl.uint16 + ) + n_hi = tl.load(KV_cache_ptr + norm_bases + 1, mask=kv_mask, other=0).to( + tl.uint16 + ) + vec_norms = (n_lo | (n_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) + + scores = vec_norms * term1 * ATTN_SCALE + scores = tl.where(kv_mask, scores, -float("inf")) + + # ============================================================ + # ONLINE SOFTMAX UPDATE (block-level) + # ============================================================ + n_e_max = tl.maximum(tl.max(scores, 0), m_prev) + re_scale = tl.exp(m_prev - n_e_max) + p = tl.exp(scores - n_e_max) + + # ============================================================ + # VALUE LOAD + DEQUANTIZE: [BLOCK_KV, BLOCK_D] + # ============================================================ + val_bases = slot_bases + KPS + + if VQB == 3: + val_addrs0 = val_bases[:, None] + val_byte_idx[None, :] + val_raw0 = tl.load( + KV_cache_ptr + val_addrs0, + mask=kv_mask[:, None] & d_mask[None, :], + other=0, + ).to(tl.int32) + val_raw1 = tl.load( + KV_cache_ptr + val_addrs0 + 1, + mask=kv_mask[:, None] & d_mask[None, :], + other=0, + ).to(tl.int32) + raw16 = val_raw0 | (val_raw1 << 8) + v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32) + + sc_bases = val_bases + VAL_DATA_BYTES + sc_lo = tl.load(KV_cache_ptr + sc_bases, mask=kv_mask, other=0).to( + tl.uint16 + ) + sc_hi = tl.load(KV_cache_ptr + sc_bases + 1, mask=kv_mask, other=0).to( + tl.uint16 + ) + v_scales = ( + (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) + ) + zr_lo = tl.load(KV_cache_ptr + sc_bases + 2, mask=kv_mask, other=0).to( + tl.uint16 + ) + zr_hi = tl.load(KV_cache_ptr + sc_bases + 3, mask=kv_mask, other=0).to( + tl.uint16 + ) + v_zeros = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) + values = v_idx * v_scales[:, None] + v_zeros[:, None] + else: # VQB == 4 + vb_idx = d_offs // 2 + vb_shift = (d_offs % 2) * 4 + val_addrs = val_bases[:, None] + vb_idx[None, :] + val_raw = tl.load( + KV_cache_ptr + val_addrs, + mask=kv_mask[:, None] & d_mask[None, :], + other=0, + ).to(tl.int32) + v_idx = ((val_raw >> vb_shift[None, :]) & 0xF).to(tl.float32) + + sc_bases = val_bases + VAL_DATA_BYTES + sc_lo = tl.load(KV_cache_ptr + sc_bases, mask=kv_mask, other=0).to( + tl.uint16 + ) + sc_hi = tl.load(KV_cache_ptr + sc_bases + 1, mask=kv_mask, other=0).to( + tl.uint16 + ) + v_scales = ( + (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) + ) + zr_lo = tl.load(KV_cache_ptr + sc_bases + 2, mask=kv_mask, other=0).to( + tl.uint16 + ) + zr_hi = tl.load(KV_cache_ptr + sc_bases + 3, mask=kv_mask, other=0).to( + tl.uint16 + ) + v_zeros = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) + values = v_idx * v_scales[:, None] + v_zeros[:, None] + + # ============================================================ + # WEIGHTED VALUE ACCUMULATION + # ============================================================ + acc = acc * re_scale + tl.sum(p[:, None] * values, 0) + l_prev = l_prev * re_scale + tl.sum(p, 0) + m_prev = n_e_max + + # Store partial result + out_base = bid * stride_mid_b + hid * stride_mid_h + sid * stride_mid_s + safe_l = tl.where(l_prev > 0.0, l_prev, 1.0) + tl.store(Mid_o_ptr + out_base + d_offs, acc / safe_l, mask=d_mask) + lse = m_prev + tl.log(safe_l) + tl.store(Mid_o_ptr + out_base + HEAD_DIM, lse) + + +# --------------------------------------------------------------------------- +# Pre-dequant kernel: Bulk dequant K (MSE+norms) and V to fp16 +# --------------------------------------------------------------------------- + + +@triton.jit +def _tq_full_dequant_kv( + KV_cache_ptr, + Block_table_ptr, + Centroids_ptr, + K_out_ptr, # [B, Hk, max_seq, D] float16 + V_out_ptr, # [B, Hk, max_seq, D] float16 + stride_ko_b, + stride_ko_h, + stride_ko_s, + stride_vo_b, + stride_vo_h, + stride_vo_s, + stride_cache_block, + stride_cache_pos, + stride_cache_head, + stride_bt_b, + HEAD_DIM: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + MSE_BYTES: tl.constexpr, + KPS: tl.constexpr, + VQB: tl.constexpr, + VAL_DATA_BYTES: tl.constexpr, + MSE_BITS: tl.constexpr, + KEY_FP8: tl.constexpr, + BLOCK_D: tl.constexpr, + NORM_CORRECTION: tl.constexpr = 0, + FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+) +): + """Full dequant: reconstruct K (MSE centroids * norm or FP8) and V to fp16.""" + pos = tl.program_id(0) + bh = tl.program_id(1) + bid = bh // NUM_KV_HEADS + hid = bh % NUM_KV_HEADS + + page_idx = pos // BLOCK_SIZE + page_off = pos % BLOCK_SIZE + block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx) + slot_base = ( + block_num * stride_cache_block + + page_off * stride_cache_pos + + hid * stride_cache_head + ) + + d_offs = tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + + # === K dequant === + ko_base = bid * stride_ko_b + hid * stride_ko_h + pos * stride_ko_s + if KEY_FP8: + k_raw = tl.load(KV_cache_ptr + slot_base + d_offs, mask=d_mask, other=0) + if FP8_E4B15: + k_recon = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32) + else: + k_recon = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32) + tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask) + else: + # MSE unpack (3-bit or 4-bit) + norms + mse_bit_off = d_offs * MSE_BITS + mse_byte_idx = mse_bit_off // 8 + mse_bit_shift = mse_bit_off % 8 + mse_umask = (1 << MSE_BITS) - 1 + + mse_raw0 = tl.load( + KV_cache_ptr + slot_base + mse_byte_idx, mask=d_mask, other=0 + ).to(tl.int32) + mse_raw1 = tl.load( + KV_cache_ptr + slot_base + mse_byte_idx + 1, mask=d_mask, other=0 + ).to(tl.int32) + raw16_key = mse_raw0 | (mse_raw1 << 8) + mse_idx = (raw16_key >> mse_bit_shift) & mse_umask + + k_mse = tl.load(Centroids_ptr + mse_idx, mask=d_mask, other=0.0) + + # Norm correction: re-normalize centroid vector to unit norm + if NORM_CORRECTION: + c_norm_sq = tl.sum(tl.where(d_mask, k_mse * k_mse, 0.0), axis=0) + c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16) + k_mse = k_mse * c_inv_norm + + # Norms at MSE_BYTES offset (no QJL bytes) + norm_base = slot_base + MSE_BYTES + n_lo = tl.load(KV_cache_ptr + norm_base).to(tl.uint16) + n_hi = tl.load(KV_cache_ptr + norm_base + 1).to(tl.uint16) + vec_norm = (n_lo | (n_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) + + k_recon = vec_norm * k_mse + tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask) + + # === V dequant === + val_base = slot_base + KPS + if VQB == 4: + vb_idx = d_offs // 2 + vb_shift = (d_offs % 2) * 4 + val_raw = tl.load(KV_cache_ptr + val_base + vb_idx, mask=d_mask, other=0).to( + tl.int32 + ) + v_idx = ((val_raw >> vb_shift) & 0xF).to(tl.float32) + + sc_base = val_base + VAL_DATA_BYTES + sc_lo = tl.load(KV_cache_ptr + sc_base).to(tl.uint16) + sc_hi = tl.load(KV_cache_ptr + sc_base + 1).to(tl.uint16) + v_scale = (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) + zr_lo = tl.load(KV_cache_ptr + sc_base + 2).to(tl.uint16) + zr_hi = tl.load(KV_cache_ptr + sc_base + 3).to(tl.uint16) + v_zero = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) + v_vals = v_idx * v_scale + v_zero + elif VQB == 3: + # 3-bit value unpack: 8 values per 3 bytes + val_bit_off = d_offs * 3 + val_byte_idx = val_bit_off // 8 + val_bit_shift = val_bit_off % 8 + val_raw0 = tl.load( + KV_cache_ptr + val_base + val_byte_idx, mask=d_mask, other=0 + ).to(tl.int32) + val_raw1 = tl.load( + KV_cache_ptr + val_base + val_byte_idx + 1, mask=d_mask, other=0 + ).to(tl.int32) + raw16_val = val_raw0 | (val_raw1 << 8) + v_idx = ((raw16_val >> val_bit_shift) & 0x7).to(tl.float32) + + sc_base = val_base + VAL_DATA_BYTES + sc_lo = tl.load(KV_cache_ptr + sc_base).to(tl.uint16) + sc_hi = tl.load(KV_cache_ptr + sc_base + 1).to(tl.uint16) + v_scale = (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) + zr_lo = tl.load(KV_cache_ptr + sc_base + 2).to(tl.uint16) + zr_hi = tl.load(KV_cache_ptr + sc_base + 3).to(tl.uint16) + v_zero = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) + v_vals = v_idx * v_scale + v_zero + else: + v_vals = tl.zeros([BLOCK_D], dtype=tl.float32) + + vo_base = bid * stride_vo_b + hid * stride_vo_h + pos * stride_vo_s + tl.store(V_out_ptr + vo_base + d_offs, v_vals.to(tl.float16), mask=d_mask) + + +# --------------------------------------------------------------------------- +# Stage 2: Reuse from triton_decode_attention.py +# --------------------------------------------------------------------------- + +# --------------------------------------------------------------------------- +# Launcher — cached constants + fused GEMM +# --------------------------------------------------------------------------- + +_layout_cache: dict = {} + + +def _get_layout(D, mse_bits, value_quant_bits, key_packed_size): + """Get cached layout constants.""" + key = (D, mse_bits, value_quant_bits, key_packed_size) + cfg = _layout_cache.get(key) + if cfg is None: + val_data_bytes = math.ceil(D * value_quant_bits / 8) + cfg = { + "mse_bytes": math.ceil(D * mse_bits / 8), + "val_data_bytes": val_data_bytes, + "mse_bits": mse_bits, + "n_centroids": 2**mse_bits, + "BLOCK_D": triton.next_power_of_2(D), + } + _layout_cache[key] = cfg + return cfg + + +def triton_turboquant_decode_attention( + query: torch.Tensor, # [B, Hq, D] — original query + kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8 + block_table: torch.Tensor, # [B, max_num_blocks] int32 + seq_lens: torch.Tensor, # [B] int32 + Pi: torch.Tensor, # [D, D] float32 + centroids: torch.Tensor, # [n_centroids] float32 + scale: float, + mse_bits: int, + key_packed_size: int, + value_quant_bits: int, + key_fp8: bool = False, + norm_correction: bool = False, + PiT: torch.Tensor | None = None, # [D, D] pre-computed Pi.T contiguous + # Pre-allocated buffers (optional, avoids per-call allocation) + mid_o_buf: torch.Tensor | None = None, + output_buf: torch.Tensor | None = None, + lse_buf: torch.Tensor | None = None, + buf_holder: Any = None, + max_num_kv_splits: int = 32, # fixed split count (must be constant for cudagraph) +) -> torch.Tensor: + """Launch fused TQ decode attention (Triton stage1 + stage2). + + Returns: output tensor [B, Hq, D] in query's dtype. + """ + B, Hq, D = query.shape + Hk = kv_cache.shape[2] + block_size = kv_cache.shape[1] + kv_group_size = Hq // Hk + device = query.device + + cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size) + + # Compute q_rot = q @ Pi.T (rotated query for MSE key scoring) + # FP8 path: pass query directly (float16); kernel casts inline. + # MSE path: still needs external GEMM (cuBLAS), so q_rot is float32. + if key_fp8: + q_rot = query.contiguous() + else: + q_float = query.float() + if PiT is None: + PiT = Pi.T.contiguous() + q_rot = (q_float @ PiT).contiguous() + + NUM_KV_SPLITS = max_num_kv_splits + + if ( + mid_o_buf is not None + and mid_o_buf.shape[0] >= B + and mid_o_buf.shape[2] >= NUM_KV_SPLITS + ): + mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :] + else: + mid_o = torch.empty( + B, + Hq, + NUM_KV_SPLITS, + D + 1, + dtype=torch.float32, + device=device, + ) + if buf_holder is not None: + buf_holder._tq_mid_o_buf = mid_o + + # Stage 1: split-KV tiled attention scoring + value accumulation + fp8_e4b15 = _use_fp8_e4b15(device.index or 0) + BLOCK_KV = 4 + grid = (B, Hq, NUM_KV_SPLITS) + _tq_decode_stage1[grid]( + q_rot, + kv_cache, + block_table, + seq_lens, + centroids, + mid_o, + q_rot.stride(0), + q_rot.stride(1), + kv_cache.stride(0), + kv_cache.stride(1), + kv_cache.stride(2), + block_table.stride(0), + mid_o.stride(0), + mid_o.stride(1), + mid_o.stride(2), + NUM_KV_HEADS=Hk, + HEAD_DIM=D, + BLOCK_SIZE=block_size, + NUM_KV_SPLITS=NUM_KV_SPLITS, + KV_GROUP_SIZE=kv_group_size, + MSE_BITS=mse_bits, + MSE_BYTES=cfg["mse_bytes"], + KPS=key_packed_size, + VQB=value_quant_bits, + VAL_DATA_BYTES=cfg["val_data_bytes"], + ATTN_SCALE=scale, + BLOCK_D=cfg["BLOCK_D"], + BLOCK_KV=BLOCK_KV, + KEY_FP8=1 if key_fp8 else 0, + NORM_CORRECTION=1 if norm_correction else 0, + FP8_E4B15=fp8_e4b15, + num_warps=1, + num_stages=1, + ) + + # Stage 2: Reduce across KV splits + if output_buf is not None and output_buf.shape[0] >= B: + output = output_buf[:B, :Hq, :D] + else: + output = torch.empty(B, Hq, D, dtype=torch.float32, device=device) + if buf_holder is not None: + buf_holder._tq_output_buf = output + if lse_buf is not None and lse_buf.shape[0] >= B: + lse = lse_buf[:B, :Hq] + else: + lse = torch.empty(B, Hq, dtype=torch.float32, device=device) + if buf_holder is not None: + buf_holder._tq_lse_buf = lse + + grid2 = (B, Hq) + _fwd_kernel_stage2[grid2]( + mid_o, + output, + lse, + seq_lens, + mid_o.stride(0), + mid_o.stride(1), + mid_o.stride(2), + output.stride(0), + output.stride(1), + lse.stride(0), + NUM_KV_SPLITS=NUM_KV_SPLITS, + BLOCK_DV=cfg["BLOCK_D"], + Lv=D, + num_warps=4, + num_stages=2, + ) + + return output.to(query.dtype) diff --git a/vllm/v1/attention/ops/triton_turboquant_store.py b/vllm/v1/attention/ops/triton_turboquant_store.py new file mode 100644 index 00000000000..3da3347d5df --- /dev/null +++ b/vllm/v1/attention/ops/triton_turboquant_store.py @@ -0,0 +1,441 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused Triton kernels for TurboQuant KV store. + +Two kernels: +1. _tq_fused_store_fp8: FP8 key scatter + value uniform quantization. +2. _tq_fused_store_mse: Fused binary-search bucketize + MSE index + packing + value quantization. + +The launcher `triton_turboquant_store` selects the appropriate kernel. +""" + +import math + +import torch + +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.triton_turboquant_decode import _use_fp8_e4b15 + +# ═══════════════════════════════════════════════════════════════════════ +# Shared: value uniform quantization + pack + scale/zero store +# ═══════════════════════════════════════════════════════════════════════ + + +@triton.jit +def _store_quantized_value( + Value_ptr, + KV_cache_ptr, + base, # pid * D offset into Value_ptr + slot_base, # byte offset into KV_cache_ptr for this slot+head + d_offs, # tl.arange(0, BLOCK_D) + d_mask, # d_offs < D + D: tl.constexpr, + KPS: tl.constexpr, + VQB: tl.constexpr, + VAL_DATA_BYTES: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_VAL: tl.constexpr, + BLOCK_GRP: tl.constexpr, +): + """Uniform quantization of values to VQB bits, pack, and store with scale/zero.""" + val_cache_offset = KPS + + if VQB == 3: + val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to( + tl.float32 + ) + val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0) + val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0) + v_scale = (val_max - val_min) / 7.0 + v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8) + + q_vals = tl.minimum( + tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 7 + ) + + grp_offs = tl.arange(0, BLOCK_GRP) + grp_mask = grp_offs < (D // 8) + q_grp = tl.reshape(q_vals, [BLOCK_GRP, 8]) + shifts_3bit = tl.arange(0, 8) * 3 + packed_24 = tl.sum(q_grp << shifts_3bit[None, :], axis=1) + b0 = (packed_24 & 0xFF).to(tl.uint8) + b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8) + b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8) + tl.store( + KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3, + b0, + mask=grp_mask, + ) + tl.store( + KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3 + 1, + b1, + mask=grp_mask, + ) + tl.store( + KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3 + 2, + b2, + mask=grp_mask, + ) + + sc_offset = val_cache_offset + VAL_DATA_BYTES + sc_f16 = v_scale.to(tl.float16) + sc_u16 = sc_f16.to(tl.uint16, bitcast=True) + tl.store(KV_cache_ptr + slot_base + sc_offset, (sc_u16 & 0xFF).to(tl.uint8)) + tl.store( + KV_cache_ptr + slot_base + sc_offset + 1, + ((sc_u16 >> 8) & 0xFF).to(tl.uint8), + ) + zr_f16 = val_min.to(tl.float16) + zr_u16 = zr_f16.to(tl.uint16, bitcast=True) + tl.store(KV_cache_ptr + slot_base + sc_offset + 2, (zr_u16 & 0xFF).to(tl.uint8)) + tl.store( + KV_cache_ptr + slot_base + sc_offset + 3, + ((zr_u16 >> 8) & 0xFF).to(tl.uint8), + ) + + else: # VQB == 4 + val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to( + tl.float32 + ) + val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0) + val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0) + v_scale = (val_max - val_min) / 15.0 + v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8) + + # Quantize all D elements from register (no re-load) + q_all = tl.minimum( + tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 15 + ) + # Reshape to pairs and pack two 4-bit values per byte + q_pairs = tl.reshape(q_all, [BLOCK_D // 2, 2]) + shifts_4 = tl.arange(0, 2) * 4 + packed_val = tl.sum((q_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8) + val_offs = tl.arange(0, BLOCK_D // 2) + val_mask = val_offs < VAL_DATA_BYTES + tl.store( + KV_cache_ptr + slot_base + val_cache_offset + val_offs, + packed_val, + mask=val_mask, + ) + + sc_offset = val_cache_offset + VAL_DATA_BYTES + sc_f16 = v_scale.to(tl.float16) + sc_u16 = sc_f16.to(tl.uint16, bitcast=True) + tl.store(KV_cache_ptr + slot_base + sc_offset, (sc_u16 & 0xFF).to(tl.uint8)) + tl.store( + KV_cache_ptr + slot_base + sc_offset + 1, + ((sc_u16 >> 8) & 0xFF).to(tl.uint8), + ) + zr_f16 = val_min.to(tl.float16) + zr_u16 = zr_f16.to(tl.uint16, bitcast=True) + tl.store(KV_cache_ptr + slot_base + sc_offset + 2, (zr_u16 & 0xFF).to(tl.uint8)) + tl.store( + KV_cache_ptr + slot_base + sc_offset + 3, + ((zr_u16 >> 8) & 0xFF).to(tl.uint8), + ) + + +# ═══════════════════════════════════════════════════════════════════════ +# FP8 key store + value uniform quantization +# ═══════════════════════════════════════════════════════════════════════ + + +@triton.jit +def _tq_fused_store_fp8( + Key_ptr, # [NH, D] float16/bfloat16 — raw keys + Value_ptr, # [NH, D] float16/bfloat16 — raw values + KV_cache_ptr, # [total_bytes] uint8 (flattened view) + Slot_mapping_ptr, # [N] int32 — per-token slot indices + # Cache strides (for computing byte offsets) + stride_cache_block: tl.constexpr, + stride_cache_pos: tl.constexpr, + stride_cache_head: tl.constexpr, + # Dimensions + D: tl.constexpr, + H: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, + # TQ layout + KPS: tl.constexpr, + # Value quantization + VQB: tl.constexpr, + VAL_DATA_BYTES: tl.constexpr, + # Packing block sizes + BLOCK_VAL: tl.constexpr, + BLOCK_GRP: tl.constexpr = 16, + FP8_E4B15: tl.constexpr = 0, # 1 = e4b15 (Ampere/Ada), 0 = e4nv (Hopper+) +): + """FP8 key cast+scatter + value uniform quantization.""" + pid = tl.program_id(0) + token_idx = pid // H + head_idx = pid % H + + slot = tl.load(Slot_mapping_ptr + token_idx) + if slot < 0: + return + blk = slot // BLOCK_SIZE + off = slot % BLOCK_SIZE + slot_base = ( + blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head + ) + + base = pid * D + + # ── FP8 KEY: cast to FP8 in-kernel and store ───────────────── + d_offs = tl.arange(0, BLOCK_D) + d_mask = d_offs < D + k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0) + k_fp8 = k_vals.to(tl.float8e4b15) if FP8_E4B15 else k_vals.to(tl.float8e4nv) + k_bytes = k_fp8.to(tl.uint8, bitcast=True) + tl.store(KV_cache_ptr + slot_base + d_offs, k_bytes, mask=d_mask) + + # ── VALUE QUANTIZE + PACK ─────────────────────────────────────── + _store_quantized_value( + Value_ptr, + KV_cache_ptr, + base, + slot_base, + d_offs, + d_mask, + D=D, + KPS=KPS, + VQB=VQB, + VAL_DATA_BYTES=VAL_DATA_BYTES, + BLOCK_D=BLOCK_D, + BLOCK_VAL=BLOCK_VAL, + BLOCK_GRP=BLOCK_GRP, + ) + + +# ═══════════════════════════════════════════════════════════════════════ +# Fused MSE store: bucketize + MSE index pack + norm store + value pack +# (eliminates 4 PyTorch kernel launches per layer vs pack-only kernel) +# ═══════════════════════════════════════════════════════════════════════ + + +@triton.jit +def _tq_fused_store_mse( + # Post-rotation inputs + Y_ptr, # [NH, D] float32 — rotated normalized keys (x_hat @ PiT) + Norms_ptr, # [NH] float32 — key vector norms (||k||) + Value_ptr, # [NH, D] float32 — raw values + # Quantization tables + Midpoints_ptr, # [n_centroids-1] float32 + # Cache and indexing + KV_cache_ptr, # [total_bytes] uint8 (flattened view) + Slot_mapping_ptr, # [N] int32 — per-token slot indices + # Cache strides + stride_cache_block: tl.constexpr, + stride_cache_pos: tl.constexpr, + stride_cache_head: tl.constexpr, + # Dimensions + D: tl.constexpr, + H: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, + # TQ layout + MSE_BYTES: tl.constexpr, + KPS: tl.constexpr, + # Value quantization + VQB: tl.constexpr, + VAL_DATA_BYTES: tl.constexpr, + # Packing block sizes + BLOCK_VAL: tl.constexpr, + # MSE params + MSE_BITS: tl.constexpr, + N_CENTROIDS: tl.constexpr, + BLOCK_GRP: tl.constexpr = 16, +): + """Fused MSE quantize + pack + store. + + Performs binary-search bucketize, MSE index packing, norm storage, + and value quantization in one kernel. + """ + pid = tl.program_id(0) + token_idx = pid // H + head_idx = pid % H + + slot = tl.load(Slot_mapping_ptr + token_idx) + if slot < 0: + return + blk = slot // BLOCK_SIZE + off = slot % BLOCK_SIZE + slot_base = ( + blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head + ) + + base = pid * D + d_offs = tl.arange(0, BLOCK_D) + d_mask = d_offs < D + + # ── 1. BINARY SEARCH BUCKETIZE ─────────────────────────────────── + # Midpoints are sorted (N_CENTROIDS-1 values); binary search finds + # insertion point in MSE_BITS iterations vs N_CENTROIDS-1 for linear. + y_vec = tl.load(Y_ptr + base + d_offs, mask=d_mask, other=0.0) + lo = tl.zeros([BLOCK_D], dtype=tl.int32) + hi = tl.full([BLOCK_D], N_CENTROIDS - 1, dtype=tl.int32) + for _ in range(MSE_BITS): + mid = (lo + hi) >> 1 + # Clamp to valid midpoint index [0, N_CENTROIDS-2] for load safety; + # the search result (lo) is still correct since converged lanes + # don't change. + safe_mid = tl.minimum(mid, N_CENTROIDS - 2) + mid_val = tl.load(Midpoints_ptr + safe_mid, mask=d_mask, other=0.0) + lo = tl.where(y_vec >= mid_val, mid + 1, lo) + hi = tl.where(y_vec >= mid_val, hi, mid) + idx = tl.minimum(lo, N_CENTROIDS - 1) + + # ── 2. PACK MSE INDICES from register idx ───────────────────────── + if MSE_BITS == 4: + idx_pairs = tl.reshape(idx, [BLOCK_D // 2, 2]) + shifts_4 = tl.arange(0, 2) * 4 + packed = tl.sum((idx_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8) + mse_offs = tl.arange(0, BLOCK_D // 2) + mse_mask = mse_offs < MSE_BYTES + tl.store(KV_cache_ptr + slot_base + mse_offs, packed, mask=mse_mask) + + elif MSE_BITS == 3: + grp_offs = tl.arange(0, BLOCK_GRP) + grp_mask = grp_offs < (D // 8) + idx_grp = tl.reshape(idx, [BLOCK_GRP, 8]) + shifts_3 = tl.arange(0, 8) * 3 + packed_24 = tl.sum((idx_grp & 0x7) << shifts_3[None, :], axis=1) + b0 = (packed_24 & 0xFF).to(tl.uint8) + b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8) + b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8) + tl.store(KV_cache_ptr + slot_base + grp_offs * 3, b0, mask=grp_mask) + tl.store(KV_cache_ptr + slot_base + grp_offs * 3 + 1, b1, mask=grp_mask) + tl.store(KV_cache_ptr + slot_base + grp_offs * 3 + 2, b2, mask=grp_mask) + + # ── 3. STORE vec_norm (fp16, 2 bytes) ───────────────────────────── + norm_offset = MSE_BYTES + + vn_f16 = tl.load(Norms_ptr + pid).to(tl.float16) + vn_u16 = vn_f16.to(tl.uint16, bitcast=True) + tl.store(KV_cache_ptr + slot_base + norm_offset, (vn_u16 & 0xFF).to(tl.uint8)) + tl.store( + KV_cache_ptr + slot_base + norm_offset + 1, ((vn_u16 >> 8) & 0xFF).to(tl.uint8) + ) + + # ── 4. VALUE QUANTIZE + PACK ────────────────────────────────────── + _store_quantized_value( + Value_ptr, + KV_cache_ptr, + base, + slot_base, + d_offs, + d_mask, + D=D, + KPS=KPS, + VQB=VQB, + VAL_DATA_BYTES=VAL_DATA_BYTES, + BLOCK_D=BLOCK_D, + BLOCK_VAL=BLOCK_VAL, + BLOCK_GRP=BLOCK_GRP, + ) + + +# ═══════════════════════════════════════════════════════════════════════ +# Launcher +# ═══════════════════════════════════════════════════════════════════════ + + +def triton_turboquant_store( + key: torch.Tensor, # [N, H, D] — raw keys (post-RoPE) + value: torch.Tensor, # [N, H, D] — raw values + kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8 + slot_mapping: torch.Tensor, # [N] int32 + PiT: torch.Tensor, # [D, D] float32 + midpoints: torch.Tensor, # [n_centroids-1] float32 + mse_bits: int, + key_packed_size: int, + value_quant_bits: int, + key_fp8: bool = False, +): + """Launch TQ store kernel (FP8 or MSE path).""" + N, H, D = key.shape + NH = N * H + block_size = kv_cache.shape[1] + BLOCK_D = triton.next_power_of_2(D) + mse_bytes = math.ceil(D * mse_bits / 8) + n_centroids = 2**mse_bits + + val_data_bytes = math.ceil(D * value_quant_bits / 8) + + BLOCK_VAL = triton.next_power_of_2(val_data_bytes) + + # Cache strides (element_size=1 for uint8, so stride in bytes = stride()) + stride_block = kv_cache.stride(0) + stride_pos = kv_cache.stride(1) + stride_head = kv_cache.stride(2) + + block_grp = triton.next_power_of_2(D // 8) if D >= 8 else 1 + + # ── FP8 PATH: in-kernel FP8 cast + scatter via fp8 kernel ── + if key_fp8: + k_flat = key.reshape(NH, D).contiguous() + v_flat = value.reshape(NH, D).contiguous() + + fp8_e4b15 = _use_fp8_e4b15(key.device.index or 0) + + grid = (NH,) + _tq_fused_store_fp8[grid]( + k_flat, + v_flat, + kv_cache.view(-1), + slot_mapping, + stride_cache_block=stride_block, + stride_cache_pos=stride_pos, + stride_cache_head=stride_head, + D=D, + H=H, + BLOCK_SIZE=block_size, + BLOCK_D=BLOCK_D, + KPS=key_packed_size, + VQB=value_quant_bits, + VAL_DATA_BYTES=val_data_bytes, + BLOCK_VAL=BLOCK_VAL, + BLOCK_GRP=block_grp, + FP8_E4B15=fp8_e4b15, + num_warps=4, + num_stages=1, + ) + return + + # ── MSE PATH: external GEMM + fused bucketize/pack kernel ── + # Normalize + rotation GEMM externally (cuBLAS is faster than in-kernel) + k_flat = key.float().reshape(NH, D) + norms = k_flat.norm(dim=1, keepdim=True) + x_hat = k_flat / (norms + 1e-8) + y = x_hat @ PiT + + v_flat = value.float().reshape(NH, D) + + # Fused kernel: bucketize + MSE index pack + norm store + value pack + grid = (NH,) + _tq_fused_store_mse[grid]( + y, + norms.squeeze(1), + v_flat, + midpoints, + kv_cache.view(-1), + slot_mapping, + stride_cache_block=stride_block, + stride_cache_pos=stride_pos, + stride_cache_head=stride_head, + D=D, + H=H, + BLOCK_SIZE=block_size, + BLOCK_D=BLOCK_D, + MSE_BYTES=mse_bytes, + KPS=key_packed_size, + VQB=value_quant_bits, + VAL_DATA_BYTES=val_data_bytes, + BLOCK_VAL=BLOCK_VAL, + MSE_BITS=mse_bits, + N_CENTROIDS=n_centroids, + BLOCK_GRP=block_grp, + num_warps=4, + num_stages=1, + ) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index fa5395685e6..30061462008 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import ( MLAAttentionSpec, SinkFullAttentionSpec, SlidingWindowSpec, + TQFullAttentionSpec, ) from vllm.v1.request import Request @@ -209,7 +210,7 @@ class SingleTypeKVCacheManager(ABC): cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) ) req_blocks.extend(allocated_blocks) - if type(self.kv_cache_spec) is FullAttentionSpec: + if type(self.kv_cache_spec) in (FullAttentionSpec, TQFullAttentionSpec): self.new_block_ids.extend(b.block_id for b in allocated_blocks) def allocate_new_blocks( @@ -237,7 +238,7 @@ class SingleTypeKVCacheManager(ABC): else: new_blocks = self.block_pool.get_new_blocks(num_new_blocks) req_blocks.extend(new_blocks) - if type(self.kv_cache_spec) is FullAttentionSpec: + if type(self.kv_cache_spec) in (FullAttentionSpec, TQFullAttentionSpec): self.new_block_ids.extend(b.block_id for b in new_blocks) return new_blocks @@ -1114,6 +1115,7 @@ class SinkFullAttentionManager(FullAttentionManager): spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = { FullAttentionSpec: FullAttentionManager, + TQFullAttentionSpec: FullAttentionManager, MLAAttentionSpec: FullAttentionManager, SlidingWindowSpec: SlidingWindowManager, ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager, diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 6f8ad8e7d8e..8aed95ddd0e 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -245,6 +245,32 @@ class FullAttentionSpec(AttentionSpec): ) +@dataclass(frozen=True, kw_only=True) +class TQFullAttentionSpec(FullAttentionSpec): + """FullAttentionSpec with TQ-aware page size. + + Python equivalent of the C++ TQ4FullAttentionSpec. Overrides + real_page_size_bytes to use TQ slot bytes instead of the raw + head_size * dtype formula. + """ + + tq_slot_size: int = 0 + + @property + def real_page_size_bytes(self) -> int: + if self.tq_slot_size > 0: + return self.block_size * self.num_kv_heads * self.tq_slot_size + return super().real_page_size_bytes + + @classmethod + def merge(cls, specs: list[Self]) -> Self: + merged = super().merge(specs) + assert all(s.tq_slot_size == specs[0].tq_slot_size for s in specs), ( + "All TQ layers in the same KV cache group must use the same tq_slot_size." + ) + return replace(merged, tq_slot_size=specs[0].tq_slot_size) + + @dataclass(frozen=True, kw_only=True) class MLAAttentionSpec(FullAttentionSpec): # TODO(Lucas/Chen): less hacky way to do this diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 83fc12cb5c3..5780624c222 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -120,7 +120,7 @@ class KVBlockZeroer: for group in attn_groups_iter: spec = group.kv_cache_spec - if type(spec) is not FullAttentionSpec: + if not isinstance(spec, FullAttentionSpec): continue if group.kv_cache_group_id >= len(kernel_block_sizes): continue From 3abf8584432acdd66bad723f9481f379ee1b3ad9 Mon Sep 17 00:00:00 2001 From: wliao2 Date: Tue, 14 Apr 2026 20:02:35 -0700 Subject: [PATCH 008/150] [Test] Refactor hard coded device string in test files under compile/quantization/models/model_executor folders (#38901) Signed-off-by: Liao, Wei --- tests/basic_correctness/test_cumem.py | 18 ++++++++++-------- .../passes/distributed/test_async_tp.py | 5 +++-- .../distributed/test_fusion_all_reduce.py | 6 ++++-- .../distributed/test_sequence_parallelism.py | 6 ++++-- tests/compile/passes/test_fusion_attn.py | 3 ++- .../passes/test_mla_attn_quant_fusion.py | 3 ++- tests/compile/passes/test_noop_elimination.py | 7 +++++-- .../passes/test_scatter_split_replace.py | 5 ++++- tests/compile/passes/test_split_coalescing.py | 5 ++++- tests/compile/test_config.py | 6 ++++-- tests/compile/test_graph_partition.py | 7 +++++-- tests/compile/test_rotary_embedding_compile.py | 4 +++- tests/compile/test_structured_logging.py | 4 +++- .../model_executor/test_eagle_quantization.py | 5 +++-- .../multimodal/pooling/test_intern_vit.py | 12 ++++++++---- tests/models/multimodal/pooling/test_radio.py | 12 ++++++++---- tests/models/test_utils.py | 10 ++++++++-- tests/quantization/test_fp8.py | 10 +++++++--- tests/quantization/test_per_token_kv_cache.py | 14 ++++++++------ tests/quantization/test_quark.py | 12 +++++++----- tests/quantization/test_torchao.py | 17 +++++++++-------- tests/test_config.py | 6 ++++-- tests/v1/sample/test_topk_topp_sampler.py | 2 +- tests/v1/test_tensor_ipc_queue.py | 9 ++++++--- 24 files changed, 122 insertions(+), 66 deletions(-) diff --git a/tests/basic_correctness/test_cumem.py b/tests/basic_correctness/test_cumem.py index b1a16cfcaba..8d8f87f0a3c 100644 --- a/tests/basic_correctness/test_cumem.py +++ b/tests/basic_correctness/test_cumem.py @@ -13,6 +13,8 @@ from vllm.utils.mem_constants import GiB_bytes from ..utils import create_new_process_for_each_test, requires_fp8 +DEVICE_TYPE = current_platform.device_type + @create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn") def test_python_error(): @@ -26,13 +28,13 @@ def test_python_error(): tensors = [] with allocator.use_memory_pool(): # allocate 70% of the total memory - x = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda") + x = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE) tensors.append(x) # release the memory allocator.sleep() # allocate more memory than the total memory - y = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda") + y = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE) tensors.append(y) with pytest.raises(RuntimeError): # when the allocator is woken up, it should raise an error @@ -44,17 +46,17 @@ def test_python_error(): def test_basic_cumem(): # some tensors from default memory pool shape = (1024, 1024) - x = torch.empty(shape, device="cuda") + x = torch.empty(shape, device=DEVICE_TYPE) x.zero_() # some tensors from custom memory pool allocator = CuMemAllocator.get_instance() with allocator.use_memory_pool(): # custom memory pool - y = torch.empty(shape, device="cuda") + y = torch.empty(shape, device=DEVICE_TYPE) y.zero_() y += 1 - z = torch.empty(shape, device="cuda") + z = torch.empty(shape, device=DEVICE_TYPE) z.zero_() z += 2 @@ -77,16 +79,16 @@ def test_basic_cumem(): def test_cumem_with_cudagraph(): allocator = CuMemAllocator.get_instance() with allocator.use_memory_pool(): - weight = torch.eye(1024, device="cuda") + weight = torch.eye(1024, device=DEVICE_TYPE) with allocator.use_memory_pool(tag="discard"): - cache = torch.empty(1024, 1024, device="cuda") + cache = torch.empty(1024, 1024, device=DEVICE_TYPE) def model(x): out = x @ weight cache[: out.size(0)].copy_(out) return out + 1 - x = torch.empty(128, 1024, device="cuda") + x = torch.empty(128, 1024, device=DEVICE_TYPE) # warmup model(x) diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py index 7edceee9811..4fbf958d867 100644 --- a/tests/compile/passes/distributed/test_async_tp.py +++ b/tests/compile/passes/distributed/test_async_tp.py @@ -31,6 +31,7 @@ from vllm.platforms import current_platform from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed +DEVICE_TYPE = current_platform.device_type FP8_DTYPE = current_platform.fp8_dtype() prompts = [ @@ -299,7 +300,7 @@ def async_tp_pass_on_test_model( ): set_random_seed(0) - device = torch.device(f"cuda:{local_rank}") + device = torch.device(f"{DEVICE_TYPE}:{local_rank}") torch.accelerator.set_device_index(device) torch.set_default_device(device) torch.set_default_dtype(dtype) @@ -324,7 +325,7 @@ def async_tp_pass_on_test_model( fuse_gemm_comms=True, ), ) - vllm_config.device_config = DeviceConfig(device=torch.device("cuda")) + vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) # this is a fake model name to construct the model config # in the vllm_config, it's not really used. diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index a50d3ca8e3e..e2c461e6692 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -37,6 +37,8 @@ from vllm.platforms import current_platform from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed +DEVICE_TYPE = current_platform.device_type + class TestAllReduceRMSNormModel(torch.nn.Module): def __init__( @@ -268,7 +270,7 @@ def all_reduce_fusion_pass_on_test_model( ): set_random_seed(0) - device = torch.device(f"cuda:{local_rank}") + device = torch.device(f"{DEVICE_TYPE}:{local_rank}") torch.accelerator.set_device_index(device) torch.set_default_device(device) torch.set_default_dtype(dtype) @@ -300,7 +302,7 @@ def all_reduce_fusion_pass_on_test_model( vllm_config.compilation_config.pass_config = PassConfig( fuse_allreduce_rms=True, eliminate_noops=True ) - vllm_config.device_config = DeviceConfig(device=torch.device("cuda")) + vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) vllm_config.parallel_config.rank = local_rank # Setup rank for debug path # this is a fake model name to construct the model config diff --git a/tests/compile/passes/distributed/test_sequence_parallelism.py b/tests/compile/passes/distributed/test_sequence_parallelism.py index 667ef4e04fb..7b240acead5 100644 --- a/tests/compile/passes/distributed/test_sequence_parallelism.py +++ b/tests/compile/passes/distributed/test_sequence_parallelism.py @@ -35,6 +35,8 @@ from vllm.platforms import current_platform from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed +DEVICE_TYPE = current_platform.device_type + pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") FP8_DTYPE = current_platform.fp8_dtype() @@ -228,7 +230,7 @@ def sequence_parallelism_pass_on_test_model( ): set_random_seed(0) - device = torch.device(f"cuda:{local_rank}") + device = torch.device(f"{DEVICE_TYPE}:{local_rank}") torch.accelerator.set_device_index(device) torch.set_default_device(device) torch.set_default_dtype(dtype) @@ -258,7 +260,7 @@ def sequence_parallelism_pass_on_test_model( eliminate_noops=True, ), ) # NoOp needed for fusion - device_config = DeviceConfig(device=torch.device("cuda")) + device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) # this is a fake model name to construct the model config # in the vllm_config, it's not really used. diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py index 40904f2db11..b776f6af98a 100644 --- a/tests/compile/passes/test_fusion_attn.py +++ b/tests/compile/passes/test_fusion_attn.py @@ -41,6 +41,7 @@ from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode +DEVICE_TYPE = current_platform.device_type FP8_DTYPE = current_platform.fp8_dtype() FP4_DTYPE = torch.uint8 @@ -300,7 +301,7 @@ def test_attention_quant_pattern( custom_ops_list = custom_ops.split(",") if custom_ops else [] - device = torch.device("cuda:0") + device = torch.device(f"{DEVICE_TYPE}:0") torch.set_default_dtype(dtype) torch.manual_seed(42) diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py index ce1fa642ad5..a5875a6b396 100644 --- a/tests/compile/passes/test_mla_attn_quant_fusion.py +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -45,6 +45,7 @@ from vllm.v1.kv_cache_interface import MLAAttentionSpec FP8_DTYPE = current_platform.fp8_dtype() FP4_DTYPE = torch.uint8 +DEVICE_TYPE = current_platform.device_type class MLAAttentionQuantPatternModel(torch.nn.Module): @@ -356,7 +357,7 @@ def test_mla_attention_quant_pattern( custom_ops_list = custom_ops.split(",") if custom_ops else [] - device = torch.device("cuda:0") + device = torch.device(f"{DEVICE_TYPE}:0") torch.set_default_dtype(dtype) torch.manual_seed(42) diff --git a/tests/compile/passes/test_noop_elimination.py b/tests/compile/passes/test_noop_elimination.py index 412e8056f9c..c31acfaf723 100644 --- a/tests/compile/passes/test_noop_elimination.py +++ b/tests/compile/passes/test_noop_elimination.py @@ -8,6 +8,9 @@ import vllm from tests.compile.backend import TestBackend from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @@ -17,7 +20,7 @@ from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConf ) @pytest.mark.parametrize("hidden_size", [64, 4096]) def test_noop_elimination(dtype, num_tokens, hidden_size, buffer_size): - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) torch.set_default_dtype(dtype) torch.manual_seed(1) @@ -88,7 +91,7 @@ def test_non_noop_slice_preserved(): Regression test for a bug where end=-1 was treated like an inferred dimension (reshape semantics) leading to incorrect elimination. """ - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) x = torch.randn(16, 16) class SliceModel(torch.nn.Module): diff --git a/tests/compile/passes/test_scatter_split_replace.py b/tests/compile/passes/test_scatter_split_replace.py index 65996089640..e85fd9f9efc 100644 --- a/tests/compile/passes/test_scatter_split_replace.py +++ b/tests/compile/passes/test_scatter_split_replace.py @@ -13,6 +13,9 @@ from vllm.compilation.passes.utility.scatter_split_replace import ( from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass from vllm.config import CompilationConfig, CompilationMode, VllmConfig from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type class ScatterSplitReplacementModel(nn.Module): @@ -61,7 +64,7 @@ class ScatterSplitReplacementModel(nn.Module): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_scatter_split_replace(dtype): - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) torch.set_default_dtype(dtype) torch.manual_seed(0) diff --git a/tests/compile/passes/test_split_coalescing.py b/tests/compile/passes/test_split_coalescing.py index a217a4af9f2..ab7a0be1a21 100644 --- a/tests/compile/passes/test_split_coalescing.py +++ b/tests/compile/passes/test_split_coalescing.py @@ -8,6 +8,9 @@ import vllm from tests.compile.backend import TestBackend from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type class SplitCoalescingModel(torch.nn.Module): @@ -28,7 +31,7 @@ class SplitCoalescingModel(torch.nn.Module): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_split_coalescing(dtype): - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) torch.set_default_dtype(dtype) torch.manual_seed(0) diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index c0f0dcca8d7..12518b4cfa2 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -31,6 +31,8 @@ from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher # This import automatically registers `torch.ops.silly.attention` from . import silly_attention # noqa: F401 +DEVICE_TYPE = current_platform.device_type + def test_version(): # Test the version comparison logic using the private function @@ -456,7 +458,7 @@ def test_cached_compilation_config(default_vllm_config): from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape dtype = torch.bfloat16 - device = torch.device("cuda:0") + device = torch.device(f"{DEVICE_TYPE}:0") batch_size, num_qo_heads, head_size = 8, 16, 128 # access and cache default compilation config @@ -478,7 +480,7 @@ def test_cached_compilation_config(default_vllm_config): query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) query_quant = torch.compile(query_quant) - _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda") + _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE) query = torch.randn( batch_size, num_qo_heads * head_size, dtype=dtype, device=device ) diff --git a/tests/compile/test_graph_partition.py b/tests/compile/test_graph_partition.py index 09c3bd0d10e..4cb199b5897 100644 --- a/tests/compile/test_graph_partition.py +++ b/tests/compile/test_graph_partition.py @@ -15,10 +15,13 @@ from vllm.compilation.backends import ( split_graph, ) from vllm.compilation.passes.fx_utils import find_op_nodes +from vllm.platforms import current_platform # This import automatically registers `torch.ops.silly.attention` from . import silly_attention # noqa: F401 +DEVICE_TYPE = current_platform.device_type + def test_getitem_moved_to_producer_subgraph(): """ @@ -151,7 +154,7 @@ def test_consecutive_ops_in_split(): final_result = torch.sigmoid(attn_inout) return final_result - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) # Create the traced FX graph for the model x = torch.randn(8, 4) @@ -329,7 +332,7 @@ def test_builtin_empty_only_partition_is_merged(): "Expected two builtin empty_like nodes in merged non-splitting subgraph" ) - x = torch.randn(2, 3, device="cuda") + x = torch.randn(2, 3, device=DEVICE_TYPE) output_original = gm(x) output_split = split_gm(x) assert torch.allclose(output_original, output_split), "Output mismatch after split" diff --git a/tests/compile/test_rotary_embedding_compile.py b/tests/compile/test_rotary_embedding_compile.py index 76f5382534e..69a4cc05084 100644 --- a/tests/compile/test_rotary_embedding_compile.py +++ b/tests/compile/test_rotary_embedding_compile.py @@ -16,6 +16,8 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type + @support_torch_compile class RotaryEmbeddingCompileModule(torch.nn.Module): @@ -45,7 +47,7 @@ def test_rotary_embedding_torch_compile_with_custom_op(monkeypatch): monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1") monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0") - device = "cuda" + device = DEVICE_TYPE positions = torch.arange(16, device=device) query = torch.randn(16, 32, device=device, dtype=torch.bfloat16) key = torch.randn(16, 32, device=device, dtype=torch.bfloat16) diff --git a/tests/compile/test_structured_logging.py b/tests/compile/test_structured_logging.py index 7813b7429b1..10b0ed139cf 100644 --- a/tests/compile/test_structured_logging.py +++ b/tests/compile/test_structured_logging.py @@ -17,8 +17,10 @@ from vllm.config.compilation import ( ) from vllm.config.scheduler import SchedulerConfig from vllm.forward_context import set_forward_context +from vllm.platforms import current_platform MLP_SIZE = 64 +DEVICE_TYPE = current_platform.device_type @support_torch_compile @@ -71,7 +73,7 @@ class TraceStructuredCapture: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_vllm_structured_logging_artifacts(use_fresh_inductor_cache): """Test that all expected vLLM artifacts are logged during compilation.""" - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) capture = TraceStructuredCapture() diff --git a/tests/model_executor/test_eagle_quantization.py b/tests/model_executor/test_eagle_quantization.py index 519a48cae52..481715da9cd 100644 --- a/tests/model_executor/test_eagle_quantization.py +++ b/tests/model_executor/test_eagle_quantization.py @@ -10,9 +10,10 @@ from vllm.config import LoadConfig, ModelConfig, SpeculativeConfig, VllmConfig from vllm.model_executor.models.utils import get_draft_quant_config from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type DEVICES = ( - [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)] - if current_platform.is_cuda_alike() + [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))] + if not current_platform.is_cpu() else ["cpu"] ) diff --git a/tests/models/multimodal/pooling/test_intern_vit.py b/tests/models/multimodal/pooling/test_intern_vit.py index cd457c62c0a..c3f7c81b78b 100644 --- a/tests/models/multimodal/pooling/test_intern_vit.py +++ b/tests/models/multimodal/pooling/test_intern_vit.py @@ -7,6 +7,7 @@ from huggingface_hub import snapshot_download from transformers import AutoConfig, AutoModel, CLIPImageProcessor from vllm.distributed import cleanup_dist_env_and_memory +from vllm.platforms import current_platform from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from ....conftest import ImageTestAssets @@ -15,6 +16,8 @@ from ....conftest import ImageTestAssets # dynamic_module and trust_remote_code for hf_runner DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"] +DEVICE_TYPE = current_platform.device_type + @torch.inference_mode() def run_intern_vit_test( @@ -39,9 +42,9 @@ def run_intern_vit_test( hf_model = AutoModel.from_pretrained( model, dtype=torch_dtype, trust_remote_code=True - ).to("cuda") + ).to(DEVICE_TYPE) hf_outputs_per_image = [ - hf_model(pixel_value.to("cuda")).last_hidden_state + hf_model(pixel_value.to(DEVICE_TYPE)).last_hidden_state for pixel_value in pixel_values ] @@ -53,9 +56,10 @@ def run_intern_vit_test( del hf_model cleanup_dist_env_and_memory() - vllm_model = vllm_model.to("cuda", torch_dtype) + vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype) vllm_outputs_per_image = [ - vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values + vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE)) + for pixel_value in pixel_values ] del vllm_model cleanup_dist_env_and_memory() diff --git a/tests/models/multimodal/pooling/test_radio.py b/tests/models/multimodal/pooling/test_radio.py index 86b5b1b5d1f..fcab077fbba 100644 --- a/tests/models/multimodal/pooling/test_radio.py +++ b/tests/models/multimodal/pooling/test_radio.py @@ -8,6 +8,7 @@ from transformers import AutoConfig, AutoModel, CLIPImageProcessor from vllm.distributed import cleanup_dist_env_and_memory from vllm.model_executor.models.radio import RadioModel +from vllm.platforms import current_platform from vllm.transformers_utils.configs.radio import RadioConfig from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE @@ -17,6 +18,8 @@ from ....conftest import ImageTestAssets # dynamic_module and trust_remote_code for hf_runner DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"] +DEVICE_TYPE = current_platform.device_type + @torch.inference_mode() def run_radio_test( @@ -51,7 +54,7 @@ def run_radio_test( config=hf_config, dtype=torch_dtype, trust_remote_code=True, - ).to("cuda") + ).to(DEVICE_TYPE) hf_model.eval() # A HF model has image normalization as a part of model's forward @@ -62,7 +65,7 @@ def run_radio_test( hf_model.make_preprocessor_external() hf_outputs_per_image = [ - hf_model(pixel_value.to("cuda")) for pixel_value in pixel_values + hf_model(pixel_value.to(DEVICE_TYPE)) for pixel_value in pixel_values ] vllm_config = RadioConfig( @@ -71,10 +74,11 @@ def run_radio_test( ) vllm_model = RadioModel(vllm_config) vllm_model.load_weights(hf_model.state_dict()) - vllm_model = vllm_model.to("cuda", torch_dtype) + vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype) vllm_outputs_per_image = [ - vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values + vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE)) + for pixel_value in pixel_values ] del vllm_model, hf_model cleanup_dist_env_and_memory() diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py index 3d719940e3a..8d47b443657 100644 --- a/tests/models/test_utils.py +++ b/tests/models/test_utils.py @@ -10,6 +10,8 @@ from vllm.model_executor.models.utils import ( ) from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type + class ModuleWithBatchNorm(torch.nn.Module): def __init__(self): @@ -174,8 +176,12 @@ class raise_if_cuda_sync: @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") def test_merge_multimodal_embeddings_no_sync(): - inputs_embeds = torch.zeros([5, 10], dtype=torch.bfloat16, device="cuda:0") - multimodal_embeddings = [torch.ones([3, 10], dtype=torch.bfloat16, device="cuda:0")] + inputs_embeds = torch.zeros( + [5, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0" + ) + multimodal_embeddings = [ + torch.ones([3, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0") + ] is_multimodal = torch.tensor([True, False, True, True, False], device="cpu") with raise_if_cuda_sync(): _merge_multimodal_embeddings( diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 4209d59ba28..4ba0e5d3dad 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -24,6 +24,8 @@ from vllm.model_executor.layers.quantization.fp8 import ( from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type + MODELS = [ "neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV", # The checkpoint below was removed from the HF. @@ -314,7 +316,7 @@ def test_scaled_fp8_quant(dtype) -> None: # Note that we use a shape % 4 != 0 to cover edge cases, # because scaled_fp8_quant is vectorized by 4. - x = (torch.randn(size=(11, 11), device="cuda") * 13).to(dtype) + x = (torch.randn(size=(11, 11), device=DEVICE_TYPE) * 13).to(dtype) # Dynamic quantization ref_y, inv_scale = ops.scaled_fp8_quant(x, None) @@ -338,7 +340,9 @@ def test_scaled_fp8_quant(dtype) -> None: # non-contiguous input with padding m, n, padded_stride = 975, 512, 576 - padded_tensor = (torch.randn(size=(m, padded_stride), device="cuda") * 13).to(dtype) + padded_tensor = (torch.randn(size=(m, padded_stride), device=DEVICE_TYPE) * 13).to( + dtype + ) x_nc = padded_tensor[:, :n] # shape (m, n) with stride (padded_stride, 1) assert not x_nc.is_contiguous() @@ -409,7 +413,7 @@ def test_fp8_reloading( # Set model config as model_config.dtype is required in Fp8LinearMethod. default_vllm_config.model_config = ModelConfig() - with torch.device("cuda:0"): + with torch.device(f"{DEVICE_TYPE}:0"): config = Fp8Config( is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, weight_block_size=weight_block_size, diff --git a/tests/quantization/test_per_token_kv_cache.py b/tests/quantization/test_per_token_kv_cache.py index 3e660e6b00d..254e284efb5 100644 --- a/tests/quantization/test_per_token_kv_cache.py +++ b/tests/quantization/test_per_token_kv_cache.py @@ -25,11 +25,13 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed from vllm.v1.kv_cache_interface import KVQuantMode, is_quantized_kv_cache +DEVICE_TYPE = current_platform.device_type + # Skip entire module if no CUDA/ROCm GPU available pytestmark = [ pytest.mark.skipif( - not current_platform.is_cuda_alike(), - reason="Per-token-head KV cache tests require CUDA or ROCm GPU.", + current_platform.is_cpu(), + reason="Per-token-head KV cache tests require GPU.", ), ] @@ -166,7 +168,7 @@ def test_reshape_and_cache_per_token_head( ) set_random_seed(seed) - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) num_blocks = (num_tokens + block_size - 1) // block_size + 4 @@ -260,7 +262,7 @@ def test_per_token_head_round_trip_accuracy( triton_reshape_and_cache_flash_per_token_head_quant, ) - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) set_random_seed(42) num_blocks = (num_tokens + block_size - 1) // block_size + 2 @@ -323,7 +325,7 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig): triton_reshape_and_cache_flash_per_token_head_quant, ) - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) num_tokens = 4 num_heads = 2 head_size = 64 @@ -430,7 +432,7 @@ def test_triton_unified_attention_per_token_head_scale( from vllm.utils.math_utils import next_power_of_2 from vllm.v1.attention.ops.triton_unified_attention import unified_attention - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) set_random_seed(0) num_seqs = len(seq_lens) diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index c51b3461c63..9eca6cda083 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -36,6 +36,8 @@ QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse( importlib.metadata.version("amd-quark") ) >= version.parse(QUARK_MXFP4_MIN_VERSION) +DEVICE_TYPE = current_platform.device_type + if QUARK_MXFP4_AVAILABLE: from quark.torch.export.nn.modules.realquantizer import StaticScaledRealQuantizer from quark.torch.kernel import mx as mx_kernel @@ -309,7 +311,7 @@ def test_mxfp4_fused_qdq_match_quark(float_dtype: torch.dtype, scalings: list[in torch.manual_seed(0) hidden_size = 64 * 32 - inp = (torch.rand(1, hidden_size, dtype=float_dtype, device="cuda") - 0.5) * 2 + inp = (torch.rand(1, hidden_size, dtype=float_dtype, device=DEVICE_TYPE) - 0.5) * 2 for i in range(hidden_size // 32): inp[:, i * 32 : (i + 1) * 32] = ( inp[:, i * 32 : (i + 1) * 32] * scalings[i % len(scalings)] @@ -353,15 +355,15 @@ def test_mxfp4_dequant_kernel_match_quark( reorder=False, real_quantized=True, float_dtype=float_dtype, - device="cuda", + device=DEVICE_TYPE, ) - observer = qspec.observer_cls(qspec, device="cuda") + observer = qspec.observer_cls(qspec, device=DEVICE_TYPE) hidden_size = 512 shape = (11008, hidden_size) - w = (torch.rand(shape, device="cuda", dtype=float_dtype) - 0.5) * 2 + w = (torch.rand(shape, device=DEVICE_TYPE, dtype=float_dtype) - 0.5) * 2 # Make it so that different groups have different scales. for i in range(hidden_size // 32): @@ -373,7 +375,7 @@ def test_mxfp4_dequant_kernel_match_quark( scale, _ = observer._calculate_qparams() weight_quantizer.scale = scale - w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to("cuda") + w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to(DEVICE_TYPE) weight_quantizer.maybe_convert_and_transpose_scale() scale = weight_quantizer.scale diff --git a/tests/quantization/test_torchao.py b/tests/quantization/test_torchao.py index fb794baa53f..8efc6742a2d 100644 --- a/tests/quantization/test_torchao.py +++ b/tests/quantization/test_torchao.py @@ -8,6 +8,7 @@ import torch from vllm.model_executor.model_loader import get_model_loader from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type DTYPE = ["bfloat16"] TORCHAO_AVAILABLE = importlib.util.find_spec("torchao") is not None @@ -33,7 +34,7 @@ def test_pre_quantized_model(vllm_runner): @pytest.mark.parametrize( "pt_load_map_location", [ - "cuda:0", + f"{DEVICE_TYPE}:0", # {"": "cuda"}, ], ) @@ -60,7 +61,7 @@ def test_qwenvl_int8wo_model_loading_with_params(vllm_runner): model_name=model_name, quantization="torchao", dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", enforce_eager=True, ) as llm: output = llm.generate_greedy(["The capital of France is"], max_tokens=4) @@ -81,7 +82,7 @@ def test_opt_125m_awq_int4wo_model_loading_with_params(vllm_runner): model_name=model_name, quantization="torchao", dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", ) as llm: output = llm.generate_greedy(["The capital of France is"], max_tokens=4) @@ -112,7 +113,7 @@ def test_online_quant_config_dict_json(vllm_runner, enable_pickle): with vllm_runner( model_name=model_name, dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", quantization="torchao", hf_overrides=hf_overrides, enforce_eager=True, @@ -158,7 +159,7 @@ def test_online_quant_config_file(vllm_runner): with vllm_runner( model_name=model_name, dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", quantization="torchao", hf_overrides=hf_overrides, enforce_eager=True, @@ -248,7 +249,7 @@ def test_opt_125m_module_fqn_to_config_regex_model(vllm_runner): torch._dynamo.reset() model_name = "torchao-testing/opt-125m-ModuleFqnToConfig-v1-regex-0.14.0.dev" with vllm_runner( - model_name=model_name, dtype="bfloat16", pt_load_map_location="cuda:0" + model_name=model_name, dtype="bfloat16", pt_load_map_location=f"{DEVICE_TYPE}:0" ) as llm: output = llm.generate_greedy(["The capital of France is"], max_tokens=4) @@ -278,7 +279,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel(vllm_runner, monkeypat model_name=model_name, quantization="torchao", dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", enforce_eager=True, ) as llm: @@ -357,7 +358,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel_online_quant( model_name=model_name, quantization="torchao", dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", hf_overrides=hf_overrides, enforce_eager=True, ) as llm: diff --git a/tests/test_config.py b/tests/test_config.py index 364285fbf52..41d34a6cb06 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -34,6 +34,8 @@ from vllm.config.vllm import ( ) from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type + def test_compile_config_repr_succeeds(): # setup: VllmBackend mutates the config object @@ -504,8 +506,8 @@ def test_generation_config_loading(): @pytest.mark.parametrize( "pt_load_map_location", [ - "cuda", - {"": "cuda"}, + DEVICE_TYPE, + {"": DEVICE_TYPE}, ], ) def test_load_config_pt_load_map_location(pt_load_map_location): diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 511f2668075..23f1f1c1f98 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -127,7 +127,7 @@ def test_flashinfer_sampler(): # ============================================================================= -@pytest.mark.skipif("CPU" in DEVICE_TYPE, reason="CUDA/XPU not available") +@pytest.mark.skipif("cpu" in DEVICE_TYPE, reason="CUDA/XPU not available") class TestTritonTopkTopp: """Tests for the Triton top-k/top-p kernel.""" diff --git a/tests/v1/test_tensor_ipc_queue.py b/tests/v1/test_tensor_ipc_queue.py index a3fcb97ca17..a70f5d48cc5 100644 --- a/tests/v1/test_tensor_ipc_queue.py +++ b/tests/v1/test_tensor_ipc_queue.py @@ -14,6 +14,7 @@ import pytest import torch import torch.multiprocessing as torch_mp +from vllm.platforms import current_platform from vllm.v1.engine.tensor_ipc import ( TensorIpcData, TensorIpcReceiver, @@ -21,6 +22,8 @@ from vllm.v1.engine.tensor_ipc import ( ) from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder +DEVICE_TYPE = current_platform.device_type + @pytest.fixture(scope="module", autouse=True) def setup_multiprocessing(): @@ -53,7 +56,7 @@ def encoder_process( encoder = MsgpackEncoder(oob_tensor_consumer=sender) if torch.cuda.is_available(): - device = "cuda:0" + device = f"{DEVICE_TYPE}:0" tensor = torch.randn( *tensor_data["shape"], dtype=tensor_data["dtype"], device=device ) @@ -384,7 +387,7 @@ def mixed_tensor_encoder_process( # Create only CUDA tensor for IPC (CPU will be serialized) # But actually, let's just send CUDA tensor directly - cuda_tensor = torch.randn(4, 5, device="cuda:0") + cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0") # Manually send via IPC to test the mechanism cuda_tensor_shared = cuda_tensor.share_memory_() @@ -651,7 +654,7 @@ def test_ipc_disabled_mode(): # If CUDA is available, test with CUDA tensor too if torch.cuda.is_available(): - cuda_tensor = torch.randn(4, 5, device="cuda:0") + cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0") encoded_cuda = encoder.encode({"cuda_tensor": cuda_tensor}) assert len(encoded_cuda) > 0 assert tensor_queues[0].empty(), ( From bcc2306cefa4179c548d3e638e7a22a88d281733 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Wed, 15 Apr 2026 11:43:29 +0800 Subject: [PATCH 009/150] [Bugfix] Respect VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY in prefetch offloader (#37699) Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Co-authored-by: Isotr0py --- vllm/model_executor/offloader/__init__.py | 2 ++ vllm/model_executor/offloader/base.py | 14 ++++++++++++++ vllm/model_executor/offloader/prefetch.py | 9 ++++----- vllm/model_executor/offloader/uva.py | 9 +++------ 4 files changed, 23 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/offloader/__init__.py b/vllm/model_executor/offloader/__init__.py index a6522ff7c0a..f1b49c69ef9 100644 --- a/vllm/model_executor/offloader/__init__.py +++ b/vllm/model_executor/offloader/__init__.py @@ -8,6 +8,7 @@ from vllm.model_executor.offloader.base import ( create_offloader, get_offloader, set_offloader, + should_pin_memory, ) from vllm.model_executor.offloader.prefetch import PrefetchOffloader from vllm.model_executor.offloader.uva import UVAOffloader @@ -20,4 +21,5 @@ __all__ = [ "create_offloader", "get_offloader", "set_offloader", + "should_pin_memory", ] diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index 7cb0ddfd184..b8c1b6cfa48 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -10,7 +10,9 @@ from typing import TYPE_CHECKING import torch.nn as nn +import vllm.envs as envs from vllm.logger import init_logger +from vllm.utils.platform_utils import is_pin_memory_available if TYPE_CHECKING: from vllm.config import OffloadConfig @@ -18,6 +20,18 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def should_pin_memory() -> bool: + """Check if pinned memory should be used for weight offloading. + + Combines the platform capability check with the user override env var. + On unified-memory systems (e.g. GH200) pinned memory eats into GPU + memory, so users can disable it via VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY. + """ + return ( + is_pin_memory_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY + ) + + """ class relation: diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index 5bdde8c3a18..cc04367d54c 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -20,8 +20,7 @@ import torch.nn as nn # Import prefetch_ops to register custom ops at module load time import vllm.model_executor.offloader.prefetch_ops # noqa: F401 from vllm.logger import init_logger -from vllm.model_executor.offloader.base import BaseOffloader -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.model_executor.offloader.base import BaseOffloader, should_pin_memory logger = init_logger(__name__) @@ -528,7 +527,7 @@ class _ModuleOffloader: gpu_buffer = offloader._gpu_buffer assert cpu_storage is not None, "CPU storage not initialized" assert gpu_buffer is not None, "GPU buffer not assigned" - assert not is_pin_memory_available() or cpu_storage.is_pinned(), ( + assert not should_pin_memory() or cpu_storage.is_pinned(), ( f"CPU storage for {name} is not pinned! " "non_blocking=True H2D copy from non-pinned memory " "causes stream synchronization that breaks " @@ -629,7 +628,7 @@ class _CpuParamOffloader(_BaseParamOffloader): original GPU tensor is garbage collected. """ param = self._param - pin_memory = is_pin_memory_available() + pin_memory = should_pin_memory() # Create pinned CPU storage and copy current GPU data self._cpu_storage = torch.empty_strided( @@ -666,7 +665,7 @@ class _CpuParamOffloader(_BaseParamOffloader): param = self._param if param.data.device.type == "cpu": - if is_pin_memory_available() and not param.data.is_pinned(): + if should_pin_memory() and not param.data.is_pinned(): pinned = torch.empty_strided( size=param.data.size(), stride=param.data.stride(), diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index c524e43cdda..51eb1a14fcb 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -10,9 +10,9 @@ from torch.func import functional_call import vllm.envs as envs from vllm.logger import init_logger -from vllm.model_executor.offloader.base import BaseOffloader +from vllm.model_executor.offloader.base import BaseOffloader, should_pin_memory from vllm.utils.mem_utils import format_gib -from vllm.utils.platform_utils import is_pin_memory_available, is_uva_available +from vllm.utils.platform_utils import is_uva_available from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor logger = init_logger(__name__) @@ -43,10 +43,7 @@ class UVAOffloader(BaseOffloader): self.cpu_offload_bytes = 0 self.cpu_offload_params = cpu_offload_params or set() - self.pin_memory = ( - is_pin_memory_available() - and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY - ) + self.pin_memory = should_pin_memory() self.uva_offloading = ( is_uva_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA ) From 799973af4e618190bfeae053517c0f53c9f8be96 Mon Sep 17 00:00:00 2001 From: zhanqiuhu <49648934+ZhanqiuHu@users.noreply.github.com> Date: Wed, 15 Apr 2026 02:50:23 -0400 Subject: [PATCH 010/150] [CI][NIXL] Fix PD CI breakage: pin nixl-cu{12,13} versions (#39851) Signed-off-by: ZhanqiuHu --- requirements/kv_connectors.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index c0b24d99ee1..700213feda1 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -1,3 +1,5 @@ lmcache >= 0.3.9 nixl[cu13] >= 0.7.1, < 0.10.0 # Required for disaggregated prefill +nixl-cu12 >= 0.7.1, < 0.10.0 +nixl-cu13 >= 0.7.1, < 0.10.0 mooncake-transfer-engine >= 0.3.8 From 431cea3eea16012aa60dc2cfccedbaeadd62b729 Mon Sep 17 00:00:00 2001 From: Wojciech Wais <12673503+wojciech-wais@users.noreply.github.com> Date: Wed, 15 Apr 2026 10:32:47 +0200 Subject: [PATCH 011/150] [Bugfix] Fix tool_calls Iterable consumed when debug logging is enabled (#34844) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Wojciech Wais Signed-off-by: mgoin Signed-off-by: Xinyu Chen Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Signed-off-by: Rishi Puri Signed-off-by: Jaebok Lee Signed-off-by: DarkLight1337 Signed-off-by: yuwei Signed-off-by: Artem Perevedentsev Signed-off-by: Ibrahim Arshad <38925737+ibrahim1023@users.noreply.github.com> Signed-off-by: Li Signed-off-by: chaunceyjiang Signed-off-by: Kunshang Ji Signed-off-by: Kunshang Ji Signed-off-by: R Signed-off-by: Lucas Wilkinson Signed-off-by: lkm2835 Signed-off-by: Ronen Schaffer Signed-off-by: vnadathur Signed-off-by: WorldExplored Signed-off-by: Srreyansh Sethi <107075589+WorldExplored@users.noreply.github.com> Signed-off-by: Isotr0py Signed-off-by: Elham Harirpoush Signed-off-by: Yan Ma Signed-off-by: Nick Hill Signed-off-by: jackcfwang Signed-off-by: Chendi Xue Signed-off-by: Injae Ryou Signed-off-by: Richard Zou Signed-off-by: milesial Signed-off-by: Elvir Crncevic Signed-off-by: whx-sjtu <2952154980@qq.com> Signed-off-by: Lalithnarayan C Signed-off-by: PatchouliTaisa Signed-off-by: jatseng-ai Signed-off-by: jatseng-ai Signed-off-by: Matthias Gehre Signed-off-by: xaguilar-amd Signed-off-by: rdondeti Signed-off-by: Ravitez Dondeti Signed-off-by: NickLucche Signed-off-by: Peter Nguyen Signed-off-by: wang.yuqi Signed-off-by: zhuhaoran Signed-off-by: Jee Jee Li Signed-off-by: tjtanaa Signed-off-by: Jesus Federico Signed-off-by: manu Signed-off-by: ZhanqiuHu Signed-off-by: Yifan Zong Signed-off-by: Rahul-Tuli Signed-off-by: Fynn Schmitt-Ulms Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Michael Goin Signed-off-by: Benjamin Chislett Signed-off-by: Tianyu Guo Signed-off-by: leeyongjun Signed-off-by: Ziying Tao Signed-off-by: jiang1.li Signed-off-by: Vibhav Agarwal Signed-off-by: ShubyM Signed-off-by: wzhao18 Signed-off-by: Itay Etelis Signed-off-by: EdalatiAli Signed-off-by: Andreas Karatzas Signed-off-by: r266-tech Signed-off-by: Roger Wang Signed-off-by: Martin Hickey Signed-off-by: Mark McLoughlin Signed-off-by: Animesh Jain Signed-off-by: Yongye Zhu Signed-off-by: zhxchen17 Signed-off-by: EricccYang Signed-off-by: Kaicheng Yang <53411596+EricccYang@users.noreply.github.com> Signed-off-by: baoloongmao Signed-off-by: sihao.li Signed-off-by: sfeng33 <4florafeng@gmail.com> Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Signed-off-by: Zhu, Zufang Signed-off-by: Tihomir Elek Signed-off-by: yiliu30 Signed-off-by: yewentao256 Signed-off-by: Santino Ramos Signed-off-by: haosdent Signed-off-by: JartX Signed-off-by: George-ao Signed-off-by: Yuyi Ao Signed-off-by: Tyler Michael Smith Signed-off-by: Mukesh Baphna Signed-off-by: Pedram Razavi Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Michael Goin Co-authored-by: Xinyu Chen Co-authored-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Co-authored-by: Rishi Puri Co-authored-by: zzaebok <44357534+zzaebok@users.noreply.github.com> Co-authored-by: Cyrus Leung Co-authored-by: Yuwei An Co-authored-by: yuwei Co-authored-by: Artem Perevedentsev Co-authored-by: Ibrahim Arshad <38925737+ibrahim1023@users.noreply.github.com> Co-authored-by: Chuan (Richard) Li Co-authored-by: Chauncey Co-authored-by: Kunshang Ji Co-authored-by: Ganesh R Co-authored-by: Lucas Wilkinson Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Kyungmin Lee <30465912+lkm2835@users.noreply.github.com> Co-authored-by: Ronen Schaffer Co-authored-by: Srreyansh Sethi <107075589+WorldExplored@users.noreply.github.com> Co-authored-by: vnadathur Co-authored-by: vnadathur <236933696+vnadathur@users.noreply.github.com> Co-authored-by: Isotr0py Co-authored-by: Elham Co-authored-by: Yan Ma Co-authored-by: Nick Hill Co-authored-by: Chaofan Wang Co-authored-by: Chendi.Xue Co-authored-by: Injae Ryou Co-authored-by: Richard Zou Co-authored-by: milesial Co-authored-by: Elvir Crnčević Co-authored-by: Claude Sonnet 4 Co-authored-by: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Co-authored-by: Lalithnarayan C Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Luka Govedič Co-authored-by: PatchyTIS <58251192+PatchouliTIS@users.noreply.github.com> Co-authored-by: PatchouliTaisa Co-authored-by: jatseng-ai Co-authored-by: Matthias Gehre Co-authored-by: xaguilar-amd Co-authored-by: Ravitez Dondeti Co-authored-by: Nicolò Lucchesi Co-authored-by: Peter Nguyen Co-authored-by: wang.yuqi Co-authored-by: zhrrr <43847754+izhuhaoran@users.noreply.github.com> Co-authored-by: Jee Jee Li Co-authored-by: TJian Co-authored-by: Jesus Federico <14651+jefp@users.noreply.github.com> Co-authored-by: Manu Co-authored-by: zhanqiuhu <49648934+ZhanqiuHu@users.noreply.github.com> Co-authored-by: yzong-rh Co-authored-by: Fynn Schmitt-Ulms Co-authored-by: Rahul-Tuli Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Benjamin Chislett Co-authored-by: Tianyu Guo Co-authored-by: Lee Yongjun <35302114+elwhyjay@users.noreply.github.com> Co-authored-by: z1ying <55220715+z1ying@users.noreply.github.com> Co-authored-by: Li, Jiang Co-authored-by: Vibhav Agarwal Co-authored-by: vibhav-agarwal Co-authored-by: ShubyM Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Co-authored-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Itay Etelis Co-authored-by: EdalatiAli Co-authored-by: Andreas Karatzas Co-authored-by: r266-tech Co-authored-by: Roger Wang Co-authored-by: Martin Hickey Co-authored-by: Or Ozeri Co-authored-by: Mark McLoughlin Co-authored-by: Le Yang <562593859@qq.com> Co-authored-by: Animesh Jain Co-authored-by: Yongye Zhu Co-authored-by: Zhengxu Chen Co-authored-by: Kaicheng Yang <53411596+EricccYang@users.noreply.github.com> Co-authored-by: maobaolong Co-authored-by: sihao_li <165983188+1643661061leo@users.noreply.github.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> Co-authored-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Co-authored-by: zofia <110436990+zufangzhu@users.noreply.github.com> Co-authored-by: Tihomir Elek Co-authored-by: Yi Liu Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Santino Ramos <51103228+santiramos27@users.noreply.github.com> Co-authored-by: haosdent Co-authored-by: JartX Co-authored-by: Yuyi Ao Co-authored-by: Tyler Michael Smith Co-authored-by: mukesh-hai Co-authored-by: Pedram Razavi --- .../openai/test_tool_calls_serialization.py | 150 ++++++++++++++++++ vllm/entrypoints/chat_utils.py | 4 +- .../openai/chat_completion/protocol.py | 41 +++++ 3 files changed, 193 insertions(+), 2 deletions(-) create mode 100644 tests/entrypoints/openai/test_tool_calls_serialization.py diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py new file mode 100644 index 00000000000..cedc2575b80 --- /dev/null +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for tool_calls Iterable → list materialisation. + +Regression tests for https://github.com/vllm-project/vllm/issues/34792. + +Setting VLLM_LOGGING_LEVEL=debug caused tool calling to break for Mistral +models because: + 1. The OpenAI Python SDK types tool_calls as Iterable[...] in + ChatCompletionAssistantMessageParam. + 2. Pydantic v2, when validating from Python objects (not from raw JSON), + wraps Iterable fields in a one-shot lazy iterator. + 3. Debug logging called model_dump_json() which consumed that iterator. + 4. The Mistral tokenizer then saw empty tool_calls and raised + "ValueError: Unexpected tool call id ...". +""" + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + + +def _make_tool_call(tc_id: str, name: str, args: str) -> dict: + return { + "id": tc_id, + "type": "function", + "function": {"name": name, "arguments": args}, + } + + +def _make_request(messages: list) -> ChatCompletionRequest: + return ChatCompletionRequest( + model="test-model", + messages=messages, + ) + + +def test_tool_calls_list_preserved_after_model_dump(): + """tool_calls in assistant messages must be readable after model_dump_json. + + When the request is built from Python dicts (as in the Anthropic → OpenAI + conversion path), Pydantic v2 previously wrapped the Iterable tool_calls + in a one-shot iterator. model_dump_json() consumed it, leaving subsequent + readers (e.g. the Mistral tokenizer) with an empty sequence. + """ + tool_call = _make_tool_call("call_abc123", "get_weather", '{"city": "Paris"}') + messages = [ + {"role": "user", "content": "What is the weather in Paris?"}, + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + { + "role": "tool", + "tool_call_id": "call_abc123", + "content": '{"temperature": 20}', + }, + ] + + req = _make_request(messages) + + # Simulate debug logging: serialize the model (this was the trigger) + _ = req.model_dump_json() + + # The assistant message must still have accessible tool_calls afterwards + assistant_msg = req.messages[1] + assert isinstance(assistant_msg, dict) + tool_calls = assistant_msg.get("tool_calls") + assert tool_calls is not None, "tool_calls must not be None after model_dump_json" + assert isinstance(tool_calls, list), "tool_calls must be a list" + assert len(tool_calls) > 0, "tool_calls must not be empty after model_dump_json" + + +def test_tool_calls_from_generator_are_materialised(): + """tool_calls passed as a generator must be converted to list on validation.""" + tool_call = _make_tool_call("call_gen1", "search", '{"query": "vllm"}') + + def tool_calls_gen(): + yield tool_call + + messages = [ + {"role": "user", "content": "Search for vllm"}, + { + "role": "assistant", + "content": None, + "tool_calls": tool_calls_gen(), # one-shot generator + }, + ] + + req = _make_request(messages) + assistant_msg = req.messages[1] + assert isinstance(assistant_msg, dict) + + # Iterate twice — must not raise or return empty on second pass + tool_calls_first = list(assistant_msg.get("tool_calls", [])) + tool_calls_second = list(assistant_msg.get("tool_calls", [])) + + assert len(tool_calls_first) == 1, "First read must return the tool call" + assert len(tool_calls_second) == 1, "Second read must also return the tool call" + + +def test_tool_calls_list_passthrough(): + """tool_calls already provided as a list must remain a list.""" + tool_call = _make_tool_call("call_list1", "calculate", '{"expr": "2+2"}') + messages = [ + {"role": "user", "content": "Calculate 2+2"}, + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + ] + + req = _make_request(messages) + assistant_msg = req.messages[1] + assert isinstance(assistant_msg, dict) + assert isinstance(assistant_msg.get("tool_calls"), list) + + +def test_messages_without_tool_calls_unaffected(): + """Messages without tool_calls must be handled correctly.""" + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello!"}, + {"role": "assistant", "content": "Hi there!"}, + ] + + req = _make_request(messages) + # None of the messages should have tool_calls injected + for msg in req.messages: + assert isinstance(msg, dict) + assert msg.get("tool_calls") is None or msg.get("tool_calls") == [] + + +@pytest.mark.parametrize("num_tool_calls", [1, 3]) +def test_multiple_tool_calls_materialised(num_tool_calls: int): + """Multiple tool calls in a single message are all preserved.""" + tool_calls = [ + _make_tool_call(f"call_{i}", f"func_{i}", f'{{"arg": {i}}}') + for i in range(num_tool_calls) + ] + messages = [ + {"role": "user", "content": "Do things"}, + {"role": "assistant", "content": None, "tool_calls": iter(tool_calls)}, + ] + + req = _make_request(messages) + assistant_msg = req.messages[1] + assert isinstance(assistant_msg, dict) + + result_tool_calls = assistant_msg.get("tool_calls") + assert isinstance(result_tool_calls, list) + assert len(result_tool_calls) == num_tool_calls + + # Verify after model_dump_json too + _ = req.model_dump_json() + assert len(assistant_msg.get("tool_calls", [])) == num_tool_calls diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index c1324d2f039..3710473560d 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -290,7 +290,7 @@ class CustomChatCompletionMessageParam(TypedDict, total=False): tool_call_id: str | None """Tool call that this message is responding to.""" - tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None + tool_calls: list[ChatCompletionMessageToolCallParam] | None """The tool calls generated by the model, such as function calls.""" reasoning: str | None @@ -321,7 +321,7 @@ class ConversationMessage(TypedDict, total=False): name: str | None """The name of the function to call""" - tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None + tool_calls: list[ChatCompletionMessageToolCallParam] | None """The tool calls generated by the model, such as function calls.""" reasoning: str | None diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 533959df609..2bc1b6e0875 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -357,6 +357,47 @@ class ChatCompletionRequest(OpenAIBaseModel): # --8<-- [end:chat-completion-extra-params] + @model_validator(mode="before") + @classmethod + def _materialize_tool_calls_before(cls, data: Any) -> Any: + """Eagerly convert tool_calls generators/iterators to lists. + + Must run before Pydantic field validation so that one-shot + generators are not consumed during union type matching of + ChatCompletionAssistantMessageParam (which types tool_calls + as Iterable[...]). + """ + if not isinstance(data, dict): + return data + messages = data.get("messages") + if not isinstance(messages, list): + return data + for msg in messages: + if not isinstance(msg, dict): + continue + tool_calls = msg.get("tool_calls") + if tool_calls is not None and not isinstance(tool_calls, list): + msg["tool_calls"] = list(tool_calls) + return data + + @model_validator(mode="after") + def _materialize_tool_calls_after(self) -> "ChatCompletionRequest": + """Convert Pydantic ValidatorIterator wrappers back to lists. + + Even after the "before" validator converts iterables to lists, + Pydantic re-wraps them in a ValidatorIterator when validating + against ChatCompletionAssistantMessageParam's Iterable[...] type. + This "after" pass materialises those wrappers so downstream code + (tokenizers, model_dump_json) always sees plain lists. + """ + for msg in self.messages: + if not isinstance(msg, dict): + continue + tool_calls = msg.get("tool_calls") + if tool_calls is not None and not isinstance(tool_calls, list): + msg["tool_calls"] = list(tool_calls) + return self + def build_chat_params( self, default_template: str | None, From 235e1f930a29f491e06fa89c2536576a71922d73 Mon Sep 17 00:00:00 2001 From: Or Ozeri Date: Wed, 15 Apr 2026 11:53:19 +0300 Subject: [PATCH 012/150] [kv_offload+HMA][3/N]: Remove block_size from KVEvents (#36644) Signed-off-by: Or Ozeri --- .../offloading_connector/test_scheduler.py | 10 +++---- tests/v1/kv_offload/test_cpu_manager.py | 26 +++++-------------- .../kv_connector/v1/offloading/scheduler.py | 2 +- vllm/v1/kv_offload/abstract.py | 1 - vllm/v1/kv_offload/cpu/manager.py | 4 --- vllm/v1/kv_offload/cpu/spec.py | 5 ---- 6 files changed, 10 insertions(+), 38 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 291d0574e50..bdc81dc1ac3 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -124,12 +124,8 @@ def test_offloading_connector(request_runner, async_scheduling: bool): return [BlockHash(str(i).encode()) for i in int_hashes] def take_events() -> Iterable[OffloadingEvent]: - yield OffloadingEvent( - keys=to_keys([1, 2, 3]), block_size=16, medium="A", removed=False - ) - yield OffloadingEvent( - keys=to_keys([4, 5, 6]), block_size=32, medium="B", removed=True - ) + yield OffloadingEvent(keys=to_keys([1, 2, 3]), medium="A", removed=False) + yield OffloadingEvent(keys=to_keys([4, 5, 6]), medium="B", removed=True) runner.manager.take_events.side_effect = take_events events = list(runner.scheduler_connector.take_events()) @@ -137,7 +133,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): event = events[0] assert isinstance(event, BlockStored) assert event.block_hashes == to_hashes([1, 2, 3]) - assert event.block_size == 16 + assert event.block_size == 0 assert event.medium == "A" assert event.token_ids == [] assert event.parent_block_hash is None diff --git a/tests/v1/kv_offload/test_cpu_manager.py b/tests/v1/kv_offload/test_cpu_manager.py index 85629c96cea..7a2ba837eca 100644 --- a/tests/v1/kv_offload/test_cpu_manager.py +++ b/tests/v1/kv_offload/test_cpu_manager.py @@ -59,7 +59,6 @@ def verify_load_output( def verify_events( events: Iterable[OffloadingEvent], - block_size: int, expected_stores: tuple[set[int], ...] = (), expected_evictions: tuple[set[int], ...] = (), ): @@ -67,7 +66,6 @@ def verify_events( evictions: list[set[OffloadKey]] = [] for event in events: assert event.medium == CPULoadStoreSpec.medium() - assert event.block_size == block_size if event.removed: evictions.append(set(event.keys)) else: @@ -98,9 +96,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): candidate to make room for [3, 4, 5] - After complete_store([2, 3, 4, 5]), block 2 must still be present. """ - block_size = 256 manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy=eviction_policy, enable_events=True, @@ -138,10 +134,9 @@ def test_cpu_manager(): """ Tests CPUOffloadingManager with lru policy. """ - # initialize a CPU backend with a capacity of 4 blocks - block_size = 256 + # initialize a CPU manager with a capacity of 4 blocks cpu_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True + num_blocks=4, cache_policy="lru", enable_events=True ) # prepare store [1, 2] @@ -163,9 +158,7 @@ def test_cpu_manager(): # complete store [1, 2] cpu_manager.complete_store(to_keys([1, 2])) - verify_events( - cpu_manager.take_events(), block_size=block_size, expected_stores=({1, 2},) - ) + verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] assert cpu_manager.lookup(to_keys([1])) == 1 @@ -184,9 +177,7 @@ def test_cpu_manager(): ) # verify eviction event - verify_events( - cpu_manager.take_events(), block_size=block_size, expected_evictions=({1},) - ) + verify_events(cpu_manager.take_events(), expected_evictions=({1},)) # prepare store with no space assert cpu_manager.prepare_store(to_keys([1, 6])) is None @@ -241,7 +232,6 @@ def test_cpu_manager(): verify_events( cpu_manager.take_events(), - block_size=block_size, expected_stores=({3, 4, 5}, {6, 7, 8}), expected_evictions=({2, 3, 4}, {8}), ) @@ -254,7 +244,6 @@ class TestARCPolicy: self, num_blocks: int = 4, enable_events: bool = True ) -> tuple[CPUOffloadingManager, ARCCachePolicy]: manager = CPUOffloadingManager( - block_size=256, num_blocks=num_blocks, cache_policy="arc", enable_events=enable_events, @@ -289,9 +278,7 @@ class TestARCPolicy: # complete store [1, 2] cpu_manager.complete_store(to_keys([1, 2])) - verify_events( - cpu_manager.take_events(), block_size=256, expected_stores=({1, 2},) - ) + verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] assert cpu_manager.lookup(to_keys([1])) == 1 @@ -547,9 +534,8 @@ def test_filter_reused_manager(): """ Tests FilterReusedOffloadingManager with a CPUOffloadingManager. """ - block_size = 256 lru_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True + num_blocks=4, cache_policy="lru", enable_events=True ) manager = FilterReusedOffloadingManager( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 52e3ce09eb8..9fd0bed8d3e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -424,7 +424,7 @@ class OffloadingConnectorScheduler: parent_block_hash=None, token_ids=[], lora_id=None, - block_size=event.block_size, + block_size=0, medium=event.medium, lora_name=None, ) diff --git a/vllm/v1/kv_offload/abstract.py b/vllm/v1/kv_offload/abstract.py index 5b3278403a0..efa617b80f9 100644 --- a/vllm/v1/kv_offload/abstract.py +++ b/vllm/v1/kv_offload/abstract.py @@ -79,7 +79,6 @@ class PrepareStoreOutput: @dataclass class OffloadingEvent: keys: list[OffloadKey] - block_size: int medium: str # True if blocks are removed, False if stored removed: bool diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 05f97df5b10..2befa91c77a 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -33,12 +33,10 @@ class CPUOffloadingManager(OffloadingManager): def __init__( self, - block_size: int, num_blocks: int, cache_policy: Literal["lru", "arc"] = "lru", enable_events: bool = False, ): - self.block_size: int = block_size self.medium: str = CPULoadStoreSpec.medium() self._num_blocks: int = num_blocks self._num_allocated_blocks: int = 0 @@ -145,7 +143,6 @@ class CPUOffloadingManager(OffloadingManager): self.events.append( OffloadingEvent( keys=to_evict, - block_size=self.block_size, medium=self.medium, removed=True, ) @@ -188,7 +185,6 @@ class CPUOffloadingManager(OffloadingManager): self.events.append( OffloadingEvent( keys=stored_keys, - block_size=self.block_size, medium=self.medium, removed=False, ) diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 4feae8cf7d5..a0fafdeb118 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -60,12 +60,7 @@ class CPUOffloadingSpec(OffloadingSpec): kv_events_config is not None and kv_events_config.enable_kv_cache_events ) - assert len(self.gpu_block_size) == 1 - gpu_block_size = self.gpu_block_size[0] - offloaded_block_size = gpu_block_size * self.block_size_factor - self._manager = CPUOffloadingManager( - block_size=offloaded_block_size, num_blocks=self.num_blocks, cache_policy=self.eviction_policy, # type: ignore[arg-type] enable_events=enable_events, From 29e5d102050669d03992a2eb863ad364ea50fab2 Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Wed, 15 Apr 2026 17:09:20 +0800 Subject: [PATCH 013/150] fix online fp8 for MiniCPM models (#39862) Signed-off-by: Yan Ma --- vllm/model_executor/models/minicpmv.py | 51 ++++++++++++++------------ 1 file changed, 27 insertions(+), 24 deletions(-) diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index 79162eef3f6..cda07ea291e 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -1050,9 +1050,17 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): quant_config=quant_config, prefix=maybe_prefix(prefix, "resampler"), ) + self._resampler_moved = False self.make_empty_intermediate_tensors = self.llm.make_empty_intermediate_tensors + def _ensure_resampler_device(self) -> None: + if self._resampler_moved: + return + # Only move device, DO NOT touch dtype (fp8 quant needs its own dtype) + self.resampler.to(current_platform.device_type) + self._resampler_moved = True + def _parse_and_validate_vision_input( self, modality: str, @@ -1171,7 +1179,9 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + loaded = loader.load_weights(weights) + self._ensure_resampler_device() + return loaded def get_mm_mapping(self) -> MultiModelKeys: """ @@ -1276,9 +1286,7 @@ class MiniCPMV2_0(MiniCPMVBaseModel): prefix=prefix, ) - return resampler.to( - device=current_platform.device_type, dtype=torch.get_default_dtype() - ) + return resampler.to(dtype=torch.get_default_dtype()) def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor: pixel_values = data["pixel_values"] @@ -1359,9 +1367,7 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA): prefix=prefix, ) - return resampler.to( - device=current_platform.device_type, dtype=torch.get_default_dtype() - ) + return resampler.to(dtype=torch.get_default_dtype()) def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor: pixel_values = data["pixel_values"] @@ -1452,11 +1458,8 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA): quant_config=quant_config, prefix=prefix, ) - target_device = current_platform.device_type - target_dtype = torch.get_default_dtype() - if any(p.is_meta for p in resampler.parameters()): - return resampler.to_empty(device=target_device).to(dtype=target_dtype) - return resampler.to(device=target_device, dtype=target_dtype) + + return resampler.to(dtype=torch.get_default_dtype()) def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor: pixel_values = data["pixel_values"] @@ -1491,7 +1494,9 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"]) - return loader.load_weights(weights) + loaded = loader.load_weights(weights) + self._ensure_resampler_device() + return loaded class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA): @@ -1551,10 +1556,7 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA): quant_config=quant_config, prefix=prefix, ) - - return resampler.to( - device=current_platform.device_type, dtype=torch.get_default_dtype() - ) + return resampler.to(dtype=torch.get_default_dtype()) def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor: pixel_values = data["pixel_values"] @@ -1589,7 +1591,9 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"]) - return loader.load_weights(weights) + loaded = loader.load_weights(weights) + self._ensure_resampler_device() + return loaded class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA): @@ -1649,11 +1653,8 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA): quant_config=quant_config, prefix=prefix, ) - target_device = current_platform.device_type - target_dtype = torch.get_default_dtype() - if any(p.is_meta for p in resampler.parameters()): - return resampler.to_empty(device=target_device).to(dtype=target_dtype) - return resampler.to(device=target_device, dtype=target_dtype) + + return resampler.to(dtype=torch.get_default_dtype()) def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor: pixel_values = data["pixel_values"] @@ -1692,7 +1693,9 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"]) - return loader.load_weights(weights) + loaded = loader.load_weights(weights) + self._ensure_resampler_device() + return loaded _SUPPORT_VERSION = { From 60995c05b4ca3a26b92dfa7abed8f5db850301cc Mon Sep 17 00:00:00 2001 From: Zhenzhong Xu Date: Wed, 15 Apr 2026 18:38:31 +0800 Subject: [PATCH 014/150] [Quantization][Autoround][CPU] Add W4A16 Support (#38192) Signed-off-by: Zhenzhong1 Signed-off-by: Zhenzhong Xu --- tests/quantization/test_cpu_wna16.py | 1 + .../model_executor/layers/quantization/inc.py | 34 ++++++++++++++++++- 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/tests/quantization/test_cpu_wna16.py b/tests/quantization/test_cpu_wna16.py index 650bf714a07..5520dc1747a 100644 --- a/tests/quantization/test_cpu_wna16.py +++ b/tests/quantization/test_cpu_wna16.py @@ -12,6 +12,7 @@ MODELS = [ "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", # with g_idx "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", # without g_idx "RedHatAI/Qwen3-1.7B-quantized.w4a16", # with zp + "OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc", ] DTYPE = ["bfloat16"] diff --git a/vllm/model_executor/layers/quantization/inc.py b/vllm/model_executor/layers/quantization/inc.py index 29d73928478..4457555c076 100644 --- a/vllm/model_executor/layers/quantization/inc.py +++ b/vllm/model_executor/layers/quantization/inc.py @@ -414,6 +414,7 @@ class INCConfig(QuantizationConfig): def apply_xpu_w4a16_quant_layer(self, layer, prefix: str): weight_bits, group_size, sym = self.get_layer_config(layer, prefix) + if not self.check_quantized(weight_bits): if isinstance(layer, (LinearBase, ParallelLMHead)): return UnquantizedLinearMethod() @@ -437,6 +438,27 @@ class INCConfig(QuantizationConfig): ) return None + def apply_cpu_w4a16_quant_layer(self, layer, prefix: str): + weight_bits, group_size, sym = self.get_layer_config(layer, prefix) + if not self.check_quantized(weight_bits): + if isinstance(layer, (LinearBase, ParallelLMHead)): + return UnquantizedLinearMethod() + else: + return None + + if weight_bits != 4: + raise NotImplementedError( + f"INC on CPU only supports 4-bit quantization, " + f"got weight_bits={weight_bits}." + ) + if not sym: + raise NotImplementedError( + "INC W4A16 on CPU only supports symmetric quantization for now." + ) + if isinstance(layer, (LinearBase, ParallelLMHead)): + return self.apply_gptq_quant_layer(layer, prefix) + return None + def get_quant_method(self, layer: torch.nn.Module, prefix: str): if prefix and self.extra_config: for layer_name in self.extra_config: @@ -446,11 +468,21 @@ class INCConfig(QuantizationConfig): return UnquantizedLinearMethod() if current_platform.is_xpu(): return self.apply_xpu_w4a16_quant_layer(layer, prefix) - if "gptq" in self.packing_format or "gptq" in self.backend: + is_gptq = "gptq" in self.packing_format or "gptq" in self.backend + if current_platform.is_cpu() and is_gptq: + return self.apply_cpu_w4a16_quant_layer(layer, prefix) + if is_gptq: return self.apply_gptq_quant_layer(layer, prefix) if "awq" in self.packing_format or "awq" in self.backend: return self.apply_awq_quant_layer(layer, prefix) + raise NotImplementedError( + f"Unsupported quantization configuration for layer '{prefix}'. " + f"Platform: CPU={current_platform.is_cpu()}. " + f"Platform: XPU={current_platform.is_xpu()}. " + f"Format: {self.packing_format}, Backend: {self.backend}." + ) + @classmethod def override_quantization_method( cls, hf_quant_cfg, user_quant, hf_config=None From 68be0f853ed0cb131468e1f9062b05d8d7a4ab34 Mon Sep 17 00:00:00 2001 From: Csrayz <33659823+Csrayz@users.noreply.github.com> Date: Wed, 15 Apr 2026 19:24:17 +0800 Subject: [PATCH 015/150] [Metrics] Add request_id to FinishedRequestStats to enable correlation between metrics and requests (#39710) Enables external `StatLogger` plugins to correlate per-request metrics with request-level context. Also, this is a pre-requisite for Prometheus exemplars in #30972. Signed-off-by: Csrayz <33659823+Csrayz@users.noreply.github.com> --- tests/v1/metrics/test_stats.py | 8 ++++++++ vllm/v1/engine/output_processor.py | 1 + vllm/v1/metrics/stats.py | 3 +++ 3 files changed, 12 insertions(+) diff --git a/tests/v1/metrics/test_stats.py b/tests/v1/metrics/test_stats.py index 3d9315c1ae6..21f496ea4ae 100644 --- a/tests/v1/metrics/test_stats.py +++ b/tests/v1/metrics/test_stats.py @@ -26,6 +26,7 @@ def test_prefill_kv_computed_with_cache(): # Case 1: With prefix cache (1200 tokens cached) iteration_stats.update_from_finished_request( finish_reason=FinishReason.STOP, + request_id="test-req-001", num_prompt_tokens=10000, max_tokens_param=100, req_stats=req_stats, @@ -35,6 +36,7 @@ def test_prefill_kv_computed_with_cache(): finished_req = iteration_stats.finished_requests[0] assert finished_req.num_prompt_tokens == 10000 assert finished_req.num_cached_tokens == 1200 + assert finished_req.request_id == "test-req-001" # Verify calculation: prefill KV = prompt tokens - cached tokens prefill_kv_computed = finished_req.num_prompt_tokens - max( @@ -55,6 +57,7 @@ def test_prefill_kv_computed_no_cache(): # Case 2: No prefix cache iteration_stats.update_from_finished_request( finish_reason=FinishReason.STOP, + request_id="test-req-002", num_prompt_tokens=2000, max_tokens_param=100, req_stats=req_stats, @@ -64,6 +67,7 @@ def test_prefill_kv_computed_no_cache(): finished_req = iteration_stats.finished_requests[0] assert finished_req.num_prompt_tokens == 2000 assert finished_req.num_cached_tokens == 0 + assert finished_req.request_id == "test-req-002" # Verify calculation: prefill KV = full prompt when no cache prefill_kv_computed = finished_req.num_prompt_tokens - max( @@ -84,6 +88,7 @@ def test_prefill_kv_computed_edge_cases(): # Case 3: Negative num_cached_tokens (shouldn't happen, but handle gracefully) iteration_stats.update_from_finished_request( finish_reason=FinishReason.STOP, + request_id="test-req-003", num_prompt_tokens=100, max_tokens_param=10, req_stats=req_stats, @@ -96,11 +101,13 @@ def test_prefill_kv_computed_edge_cases(): finished_req.num_cached_tokens, 0 ) assert prefill_kv_computed == 100 # Should treat negative as 0 + assert finished_req.request_id == "test-req-003" # Case 4: All tokens cached (shouldn't happen in practice) iteration_stats2 = IterationStats() iteration_stats2.update_from_finished_request( finish_reason=FinishReason.STOP, + request_id="test-req-004", num_prompt_tokens=100, max_tokens_param=10, req_stats=req_stats, @@ -112,6 +119,7 @@ def test_prefill_kv_computed_edge_cases(): finished_req2.num_cached_tokens, 0 ) assert prefill_kv_computed2 == 0 # All cached, nothing computed + assert finished_req2.request_id == "test-req-004" def test_prompt_token_stats_all_computed(): diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 3d1d7a82db3..1ae89ae1968 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -799,6 +799,7 @@ class OutputProcessor: assert req_state.stats is not None iteration_stats.update_from_finished_request( finish_reason=finish_reason, + request_id=req_state.external_req_id, num_prompt_tokens=req_state.prompt_len, max_tokens_param=req_state.max_tokens_param, req_stats=req_state.stats, diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index 78faa73a88d..a7a5fb7a2d2 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -225,6 +225,7 @@ class FinishedRequestStats: """Stats associated with a finished request.""" finish_reason: "FinishReason" + request_id: str | None = None e2e_latency: float = 0.0 num_prompt_tokens: int = 0 num_generation_tokens: int = 0 @@ -427,6 +428,7 @@ class IterationStats: def update_from_finished_request( self, finish_reason: "FinishReason", + request_id: str, num_prompt_tokens: int, max_tokens_param: int | None, req_stats: RequestStateStats, @@ -458,6 +460,7 @@ class IterationStats: finished_req = FinishedRequestStats( finish_reason=finish_reason, + request_id=request_id, e2e_latency=e2e_latency, num_prompt_tokens=num_prompt_tokens, num_generation_tokens=req_stats.num_generation_tokens, From fc701c80588c215f84af0b745edcf4d127e276bc Mon Sep 17 00:00:00 2001 From: zofia <110436990+zufangzhu@users.noreply.github.com> Date: Wed, 15 Apr 2026 20:28:19 +0800 Subject: [PATCH 016/150] [XPU][MXFP4] add mxfp4 quant op for XPU (#39857) Signed-off-by: Zhu, Zufang --- vllm/_xpu_ops.py | 46 +++++++++++++++++++ .../layers/quantization/utils/mxfp4_utils.py | 4 ++ 2 files changed, 50 insertions(+) diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index e4d2d8b68ed..a0a321173d1 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -144,6 +144,46 @@ def _xpu_mxfp8_quantize_fake( return x.to(dtype), x_s.to(torch.float8_e8m0fnu) +def _xpu_mxfp4_quantize_impl( + x: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + MXFP4_BLOCK_SIZE = 32 + eps = 1e-10 + assert x.ndim == 2, "input must be 2-D" + assert x.shape[-1] % MXFP4_BLOCK_SIZE == 0, ( + f"last dimension {x.shape[-1]} must be divisible by group_size " + f"{MXFP4_BLOCK_SIZE}" + ) + assert x.is_contiguous(), "input groups must be contiguous" + + M, N = x.shape + + # Packed FP4 output: two nibbles per byte + x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8) + x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32) + + torch.ops._C.per_token_group_quant_mxfp4(x, x_q, x_s, MXFP4_BLOCK_SIZE, eps) + + x_q = x_q.view(torch.float4_e2m1fn_x2) + x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format) + return x_q, x_s + + +def _xpu_mxfp4_quantize_fake( + x: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + MXFP4_BLOCK_SIZE = 32 + M, N = x.shape + + # Packed FP4 output: two nibbles per byte + x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8) + x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32) + + x_q = x_q.view(torch.float4_e2m1fn_x2) + x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format) + return x_q, x_s + + # Global flag to ensure ops are registered only once _OPS_REGISTERED = False @@ -555,6 +595,12 @@ class xpu_ops: fake_impl=_xpu_mxfp8_quantize_fake, ) + direct_register_custom_op( + op_name="xpu_mxfp4_quantize", + op_func=_xpu_mxfp4_quantize_impl, + fake_impl=_xpu_mxfp4_quantize_fake, + ) + _OPS_REGISTERED = True diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 21c8aba1d56..51b7b29551d 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -162,3 +162,7 @@ try: quant_dequant_mxfp4 = torch.ops.vllm.quant_dequant_mxfp4 except AttributeError as error: raise error + + +def xpu_mxfp4_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.vllm.xpu_mxfp4_quantize(x) From db8d4a4a06ac747894534982b6e52f84eb125fab Mon Sep 17 00:00:00 2001 From: Chauncey Date: Wed, 15 Apr 2026 21:28:09 +0800 Subject: [PATCH 017/150] [BugFix][Graph] fix: handle empty sym_shape_indices in PiecewiseBackend. (#39395) Signed-off-by: chaunceyjiang --- .../test_dynamic_shapes_compilation.py | 44 +++++++++++++++++++ vllm/compilation/piecewise_backend.py | 20 ++++++--- 2 files changed, 59 insertions(+), 5 deletions(-) diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index bbd62237c5e..1775b2c9deb 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -222,3 +222,47 @@ def test_model_specialization_with_evaluate_guards( torch.randn(1, 10).cuda(), is_01_specialization=True, ) + + +@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") +def test_piecewise_backend_empty_sym_shape_indices(): + """Test that PiecewiseBackend handles empty sym_shape_indices correctly. + + When all inputs have static shapes (no torch.SymInt), sym_shape_indices + will be empty. The fix in PiecewiseBackend.__call__ handles this case + by using the first compiled range_entry. + """ + gc.collect() + torch.accelerator.empty_cache() + torch.accelerator.synchronize() + + # Use small max_model_len and max_num_batched_tokens to encourage + # static shape compilation with empty sym_shape_indices + llm = LLM( + model="Qwen/Qwen3-0.6B", + max_model_len=512, + max_num_batched_tokens=1, + compilation_config={ + "mode": CompilationMode.VLLM_COMPILE, + "dynamic_shapes_config": { + "type": DynamicShapesType.BACKED.value, + }, + }, + ) + + sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) + + # Generate with static shape inputs + output = llm.generate("Hello, my name is", sampling_params=sampling_params) + result = output[0].outputs[0].text + assert len(result) > 0, "Should generate non-empty output" + + # Generate again to verify compilation works with empty sym_shape_indices + output = llm.generate("The capital of France is", sampling_params=sampling_params) + result = output[0].outputs[0].text + assert len(result) > 0, "Should generate non-empty output on second run" + + del llm + gc.collect() + torch.accelerator.empty_cache() + torch.accelerator.synchronize() diff --git a/vllm/compilation/piecewise_backend.py b/vllm/compilation/piecewise_backend.py index e167a303de7..02a4dad5460 100644 --- a/vllm/compilation/piecewise_backend.py +++ b/vllm/compilation/piecewise_backend.py @@ -354,12 +354,22 @@ class PiecewiseBackend: return None def __call__(self, *args: Any) -> Any: - runtime_shape = args[self.sym_shape_indices[0]] - range_entry = self._find_range_for_shape(runtime_shape) + if self.sym_shape_indices: + runtime_shape = args[self.sym_shape_indices[0]] + range_entry = self._find_range_for_shape(runtime_shape) + assert range_entry is not None, ( + f"Shape: {runtime_shape} out of considered ranges: " + f"{self.compile_ranges}" + ) + else: + # All inputs have static shapes; use the only compiled range_entry + compiled_entries = [re for re in self.range_entries.values() if re.compiled] + assert len(compiled_entries) == 1, ( + f"Expected exactly one compiled range_entry for static shape " + f"compilation, but found {len(compiled_entries)}" + ) + range_entry = compiled_entries[0] - assert range_entry is not None, ( - f"Shape: {runtime_shape} out of considered ranges: {self.compile_ranges}" - ) assert range_entry.compiled, ( "All ranges should be compiled or loaded up front in " "PiecewiseBackend.__init__. " From 8b5531933a7b11965b58b21c19ea853487d3f78d Mon Sep 17 00:00:00 2001 From: danielafrimi <45691845+danielafrimi@users.noreply.github.com> Date: Wed, 15 Apr 2026 16:49:48 +0300 Subject: [PATCH 018/150] FIX: support language_model.backbone naming in NemotronH Nano VL quantization config (#39901) Signed-off-by: <> Co-authored-by: root --- vllm/model_executor/models/nano_nemotron_vl.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 81b3ac26026..b0424675943 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -37,6 +37,7 @@ from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM from vllm.model_executor.models.parakeet import ParakeetExtractor, ProjectedParakeet from vllm.model_executor.models.radio import RadioModel, calc_seq_lens from vllm.model_executor.models.utils import ( + WeightsMapper, init_vllm_registered_model, maybe_prefix, ) @@ -903,6 +904,12 @@ class NemotronH_Nano_VL_V2( requires_sequential_video_encoding = True """Temporarily needed for dynamic res video w/ conv3d, doesn't support bs>1 yet""" + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "language_model.backbone": "language_model.model", + }, + ) + @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: if modality.startswith("image"): From 3beb57a238b82fe90e8b99e009c876343b9d9703 Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Wed, 15 Apr 2026 21:52:58 +0800 Subject: [PATCH 019/150] [XPU] properly handle q_descale on XPU as quant query input not supported (#39676) Signed-off-by: Yan Ma Co-authored-by: Kunshang Ji --- vllm/v1/attention/backends/flash_attn.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 0e926a99b8c..6af0fa7c496 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -1031,7 +1031,9 @@ class FlashAttentionImpl(AttentionImpl): window_size=sliding_window_size, softcap=self.logits_soft_cap, fa_version=self.vllm_flash_attn_version, - q_descale=layer._q_scale.expand(descale_shape), + q_descale=layer._q_scale.expand(descale_shape) + if self.supports_quant_query_input + else None, k_descale=layer._k_scale.expand(descale_shape), v_descale=layer._v_scale.expand(descale_shape), num_splits=1 if self.batch_invariant_enabled else 0, From 3cc328a4be4976f75ce016f60bc55beee4701d1b Mon Sep 17 00:00:00 2001 From: Talor Abramovich Date: Wed, 15 Apr 2026 19:00:07 +0300 Subject: [PATCH 020/150] [SpecDecode][Benchmark] Add SPEED-bench support to benchmarking CLI (#36029) Signed-off-by: talora Co-authored-by: Benjamin Chislett --- docs/benchmarking/cli.md | 64 ++++++++++++++++++++ vllm/benchmarks/datasets/datasets.py | 88 ++++++++++++++++++++++++++++ 2 files changed, 152 insertions(+) diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index f78ae8a9536..a7e50c907d5 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -37,6 +37,7 @@ th { | HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` | | HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech` | | Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` | +| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` | | Custom | ✅ | ✅ | Local file: `data.jsonl` | | Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` | @@ -239,6 +240,69 @@ vllm bench serve \ --spec-bench-category "summarization" ``` +#### SPEED-Bench Benchmark with Speculative Decoding + +[SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) is a unified and diverse dataset for speculative decoding, supporting acceptance rate and length measurements using the Qualitative split and throughput measurements using the Throughput splits in 5 configuration of input sequence length (1k, 2k, 8k, 16k, 32k). + +!!! note + This dataset is governed by the [NVIDIA Evaluation Dataset License Agreement](https://huggingface.co/datasets/nvidia/SPEED-Bench/blob/main/License.pdf). For each dataset a user elects to use, the user is responsible for checking if the dataset license is fit for the intended purpose. The `prepare.py` script automatically fetches data from all the source datasets. + +First, download the dataset to a folder, using this one liner: + +```bash +curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 - +``` + +The command supports also the following arguments: + +- `--config`: download only a subset of the dataset: `qualitative`, `throughput_1k`, `throughput_2k`, `throughput_8k`, `throughput_16k` and `throughput_32k`. By default, it will download all subsets. +- `--output_dir`: download to a specified folder. By default, it will download to the current directory. + +Start a server with speculative decoding: + +```bash +vllm serve meta-llama/Llama-3.3-70B-Instruct \ + --speculative-config $'{"method": "eagle3", + "num_speculative_tokens": 3, + "model": "nvidia/Llama-3.3-70B-Instruct-Eagle3"}' +``` + +Run all categories in the Qualitative split: + +```bash +vllm bench serve \ + --model meta-llama/Llama-3.3-70B-Instruct \ + --dataset-name speed_bench \ + --dataset-path "/data/speed_bench" \ + --num-prompts -1 +``` + +Available categories include `[writing, roleplay, reasoning, math, coding, stem, humanities, multilingual, summarization, qa, rag]`. + +Run only a specific category like "multilingual": + +```bash +vllm bench serve \ + --model meta-llama/Llama-3.3-70B-Instruct \ + --dataset-name speed_bench \ + --dataset-path "/data/speed_bench" \ + --num-prompts -1 + --speed-bench-category "multilingual" +``` + +Run all categories in the Throughput split (2k ISL): + +```bash +vllm bench serve \ + --model meta-llama/Llama-3.3-70B-Instruct \ + --dataset-name speed_bench \ + --speed-bench-dataset-subset throughput_2k + --dataset-path "/data/speed_bench/" \ + --num-prompts -1 +``` + +Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card. + #### Other HuggingFaceDataset Examples ```bash diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index d7ba9d8787a..986c0c85625 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -25,6 +25,7 @@ from contextlib import suppress from dataclasses import dataclass, replace from functools import cache from io import BytesIO +from pathlib import Path from tempfile import NamedTemporaryFile from typing import Any, cast @@ -1422,6 +1423,7 @@ def add_dataset_parser(parser: FlexibleArgumentParser): "custom_mm", "prefix_repetition", "spec_bench", + "speed_bench", ], help="Name of the dataset to benchmark on.", ) @@ -1606,6 +1608,34 @@ def add_dataset_parser(parser: FlexibleArgumentParser): "repetition dataset.", ) + speed_bench_group = parser.add_argument_group("speed bench dataset options") + speed_bench_group.add_argument( + "--speed-bench-dataset-subset", + type=str, + default="qualitative", + choices={ + "qualitative", + "throughput_1k", + "throughput_2k", + "throughput_8k", + "throughput_16k", + "throughput_32k", + }, + help="Subset of the SPEED-Bench dataset.", + ) + speed_bench_group.add_argument( + "--speed-bench-output-len", + type=int, + default=4096, + help="Num of output tokens per request, used only for speed bench dataset.", + ) + speed_bench_group.add_argument( + "--speed-bench-category", + type=str, + default=None, + help="Category for speed bench dataset. If None, use all categories.", + ) + def add_random_dataset_base_args( parser_or_group: FlexibleArgumentParser | argparse._ArgumentGroup, @@ -2074,6 +2104,19 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: request_id_prefix=args.request_id_prefix, no_oversample=args.no_oversample, ), + "speed_bench": lambda: SpeedBench( + dataset_path=args.dataset_path, + dataset_subset=args.speed_bench_dataset_subset, + category=args.speed_bench_category, + disable_shuffle=args.disable_shuffle, + ).sample( + num_requests=args.num_prompts, + tokenizer=tokenizer, + output_len=args.speed_bench_output_len, + enable_multimodal_chat=args.enable_multimodal_chat, + request_id_prefix=args.request_id_prefix, + no_oversample=args.no_oversample, + ), } try: @@ -3551,3 +3594,48 @@ class MMStarDataset(HuggingFaceDataset): sampled_requests, num_requests, request_id_prefix, no_oversample ) return sampled_requests + + +# ----------------------------------------------------------------------------- +# Speed Bench Dataset Implementation +# ----------------------------------------------------------------------------- + + +class SpeedBench(CustomDataset): + """ + Implements the SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench + Download the dataset using: + curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py | python3 - + """ # noqa: E501 + + def __init__(self, **kwargs) -> None: + self.dataset_subset = kwargs.pop("dataset_subset", "qualitative") + self.category = kwargs.pop("category", None) + super().__init__(**kwargs) + self.load_data() + + def load_data(self) -> None: + if self.dataset_path is None: + raise ValueError("dataset_path must be provided for loading data.") + + self.data = [] + + # Load the JSONL file + jsonl_data = pd.read_json( + path_or_buf=Path(self.dataset_path) / f"{self.dataset_subset}.jsonl", + lines=True, + ) + + # check if the JSONL file has a 'turns' column + if "messages" not in jsonl_data.columns: + raise ValueError("JSONL file must contain a 'messages' column.") + + for _, row in jsonl_data.iterrows(): + # sample only from a specific category if specified + if (not self.category) or (self.category == row["category"]): + prompt = row["messages"][0]["content"] + self.data.append({"prompt": prompt}) + + random.seed(self.random_seed) + if not getattr(self, "disable_shuffle", False): + random.shuffle(self.data) From ed333105520c9610daa17cfe6be21383513b9c34 Mon Sep 17 00:00:00 2001 From: Roy Huang Date: Wed, 15 Apr 2026 09:10:49 -0700 Subject: [PATCH 021/150] [KVConnector][LMCache] Propagate cache_salt through MP connector for per-user cache isolation (#39837) Signed-off-by: royyhuang Signed-off-by: royyhuang --- .../kv_connector/v1/lmcache_mp_connector.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py index 6686b60df4a..c6d46b49af5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py @@ -215,8 +215,11 @@ class LMCacheMPRequestTracker: # Main state state: LMCacheMPRequestState = LMCacheMPRequestState.PREFETCHING + cache_salt: str = "" + def __init__(self, request: "Request"): self.request_id = request.request_id + self.cache_salt: str = request.cache_salt or "" self.all_token_ids = request.all_token_ids self.block_hashes = ConstantList(request.block_hashes) self.allocated_block_ids = [] @@ -289,6 +292,7 @@ class LMCacheMPRequestMetadata: request_id: str direction: Literal["STORE", "RETRIEVE"] op: LoadStoreOp + cache_salt: str = "" @staticmethod def GetStoreMetadata( @@ -355,6 +359,7 @@ class LMCacheMPRequestMetadata: request_id=tracker.request_id, direction="STORE", op=op, + cache_salt=tracker.cache_salt, ) # Update the request tracker @@ -421,6 +426,7 @@ class LMCacheMPRequestMetadata: request_id=tracker.request_id, direction="RETRIEVE", op=op, + cache_salt=tracker.cache_salt, ) return ret @@ -569,12 +575,14 @@ class LMCacheMPConnector(KVConnectorBase_V1): request_ids = [] ops = [] + cache_salts = [] for meta in metadata.requests: if meta.direction != "RETRIEVE": continue request_ids.append(meta.request_id) ops.append(meta.op) + cache_salts.append(meta.cache_salt) if len(request_ids) == 0: return @@ -583,7 +591,9 @@ class LMCacheMPConnector(KVConnectorBase_V1): event = torch.cuda.Event(interprocess=True) event.record() - self.worker_adapter.batched_submit_retrieve_requests(request_ids, ops, event) + self.worker_adapter.batched_submit_retrieve_requests( + request_ids, ops, event, cache_salts=cache_salts + ) def wait_for_layer_load(self, layer_name: str) -> None: """ @@ -640,11 +650,13 @@ class LMCacheMPConnector(KVConnectorBase_V1): request_ids = [] ops = [] + cache_salts = [] for meta in metadata.requests: if meta.direction != "STORE": continue request_ids.append(meta.request_id) ops.append(meta.op) + cache_salts.append(meta.cache_salt) if len(request_ids) == 0: return @@ -653,7 +665,9 @@ class LMCacheMPConnector(KVConnectorBase_V1): event = torch.cuda.Event(interprocess=True) event.record() - self.worker_adapter.batched_submit_store_requests(request_ids, ops, event) + self.worker_adapter.batched_submit_store_requests( + request_ids, ops, event, cache_salts=cache_salts + ) def get_finished( self, finished_req_ids: set[str] @@ -755,6 +769,7 @@ class LMCacheMPConnector(KVConnectorBase_V1): self.scheduler_adapter.maybe_submit_lookup_request( request.request_id, token_ids=list(request.all_token_ids), + cache_salt=tracker.cache_salt, ) ret = self.scheduler_adapter.check_lookup_result(request.request_id) From f2145efcb6dc7b5ff0ec0a321508ec7f313c6f83 Mon Sep 17 00:00:00 2001 From: daniebrill <50454544+daniebrill@users.noreply.github.com> Date: Wed, 15 Apr 2026 18:15:01 +0200 Subject: [PATCH 022/150] [BugFix] KeyError on scope["method"] for realtime api websocket in AuthenticationMiddleware (#36934) Signed-off-by: daniebrill <50454544+daniebrill@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- vllm/entrypoints/openai/server_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/openai/server_utils.py b/vllm/entrypoints/openai/server_utils.py index 02b8c335262..f8fe3366b06 100644 --- a/vllm/entrypoints/openai/server_utils.py +++ b/vllm/entrypoints/openai/server_utils.py @@ -69,7 +69,10 @@ class AuthenticationMiddleware: return token_match def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: - if scope["type"] not in ("http", "websocket") or scope["method"] == "OPTIONS": + if ( + scope["type"] not in ("http", "websocket") + or scope.get("method") == "OPTIONS" + ): # scope["type"] can be "lifespan" or "startup" for example, # in which case we don't need to do anything return self.app(scope, receive, send) From 8ad6ff0037e825e33bf31d902afcb12221ea39d1 Mon Sep 17 00:00:00 2001 From: Mark McLoughlin Date: Wed, 15 Apr 2026 17:22:20 +0100 Subject: [PATCH 023/150] [Test] Fix @create_new_process_for_each_test("fork") in interactive shell pipeline (#29130) Signed-off-by: Mark McLoughlin --- tests/utils.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/tests/utils.py b/tests/utils.py index 7af72cb730b..cfa5f525f5d 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1261,9 +1261,6 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non @functools.wraps(func) def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None: - # Make the process the leader of its own process group - # to avoid sending SIGTERM to the parent process - os.setpgrp() from _pytest.outcomes import Skipped # Create a unique temporary file to store exception info from child @@ -1283,6 +1280,9 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non pid = os.fork() print(f"Fork a new process to run a test {pid}") if pid == 0: + # Make the child process the leader of its own process group + # to avoid sending SIGTERM to the parent process + os.setpgrp() # Parent process responsible for deleting, don't delete # in child. delete_after.pop_all() @@ -1322,14 +1322,12 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non else: os._exit(0) else: - pgid = os.getpgid(pid) + # After setpgrp(), the child's pgid equals its pid + pgid = pid _pid, _exitcode = os.waitpid(pid, 0) - # ignore SIGTERM signal itself - old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN) - # kill all child processes - os.killpg(pgid, signal.SIGTERM) - # restore the signal handler - signal.signal(signal.SIGTERM, old_signal_handler) + # kill all child processes - but they may already have exited cleanly + with contextlib.suppress(ProcessLookupError): + os.killpg(pgid, signal.SIGTERM) if _exitcode != 0: # Try to read the exception from the child process exc_info = {} From 21e5a9f48e773e36e916bc8d10c4ab4aed3887a7 Mon Sep 17 00:00:00 2001 From: Monishver <79901686+Monishver11@users.noreply.github.com> Date: Wed, 15 Apr 2026 13:48:12 -0400 Subject: [PATCH 024/150] Bug/test eagle dp v2 (#39838) Signed-off-by: Monishver Chandrasekaran --- .buildkite/test_areas/distributed.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 1c9ba95183c..56e528ffb37 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -196,7 +196,6 @@ steps: - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - pytest -v -s tests/v1/distributed/test_dbo.py - - TP_SIZE=1 DP_SIZE=2 pytest -v -s tests/v1/distributed/test_eagle_dp.py - label: Distributed Tests (2 GPUs)(B200) device: b200 From 55e1a8e1035bddb0b5b63f9ddecc8b4e16fc3ef6 Mon Sep 17 00:00:00 2001 From: Zhewen Li Date: Wed, 15 Apr 2026 11:36:47 -0700 Subject: [PATCH 025/150] [Mooncake] Fix mixed MLA+Eagle block-size validation (#39596) Signed-off-by: Zhewen Li Co-authored-by: Zhewen Li Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../unit/test_mooncake_connector.py | 45 +++++++++++++++++++ .../v1/mooncake/mooncake_connector.py | 23 ++++++++-- 2 files changed, 65 insertions(+), 3 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index f21f8ecdc5c..7b6fe3af0ce 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -609,6 +609,51 @@ def test_register_kv_caches(): assert bl == tensor1[0].nbytes // tensor1.shape[1] +def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes(): + """Mixed MLA+Eagle caches should register by byte length, not shape.""" + + vllm_config = create_vllm_config( + kv_connector="MooncakeConnector", kv_role="kv_consumer" + ) + + with ( + set_current_vllm_config(vllm_config), + patch_worker_dependencies(), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Event" + ), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Thread" + ) as mock_thread, + ): + connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER) + worker = connector.connector_worker + mock_thread.return_value.is_alive.return_value = False + + worker.use_mla = True + worker.kv_topo.is_mla = True + + # MLA cache tensor: shape[-2] is the block size. + mla_cache = torch.zeros((2, 16, 96), dtype=torch.float16) + # Eagle3/GQA-like cache tensor: shape[-2] is num_kv_heads, not block size. + eagle_cache = torch.zeros((2, 16, 8, 64), dtype=torch.float16) + kv_caches = {"mla_layer": mla_cache, "eagle_layer": eagle_cache} + + with patch.object( + worker.engine, "batch_register_memory", return_value=0 + ) as mock_batch_register: + connector.register_kv_caches(kv_caches) + + mock_batch_register.assert_called_once() + registered_ptrs, registered_lens = mock_batch_register.call_args[0] + assert registered_ptrs == [mla_cache.data_ptr(), eagle_cache.data_ptr()] + assert registered_lens == [mla_cache.nbytes, eagle_cache.nbytes] + assert worker.block_len_per_layer == [ + mla_cache.nbytes // mla_cache.shape[0], + eagle_cache.nbytes // eagle_cache.shape[0], + ] + + @pytest.mark.asyncio @patch( "vllm.distributed.kv_transfer.kv_connector.v1.mooncake." diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 6c17bf82ace..67603e10ff6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import ( EngineId, TpKVTopology, get_current_attn_backend, + get_current_attn_backends, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, @@ -47,6 +48,7 @@ from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.request import RequestStatus +from vllm.v1.worker.utils import select_common_block_size logger = init_logger(__name__) @@ -751,6 +753,7 @@ class MooncakeConnectorWorker: self.model_config = vllm_config.model_config self.cache_config = vllm_config.cache_config self.use_mla = self.model_config.use_mla + self._sync_block_size_with_kernel() # Get the attention backend from the first layer # NOTE (NickLucche) models with multiple backends are not supported yet @@ -777,6 +780,23 @@ class MooncakeConnectorWorker: self._xfer_meta_decoder = msgspec.msgpack.Decoder(MooncakeXferMetadata) self._xfer_resp_decoder = msgspec.msgpack.Decoder(MooncakeXferResponse) + def _sync_block_size_with_kernel(self) -> None: + # When speculative decoding (e.g. Eagle) is enabled, the main model + # and draft model may use different attention backends with different + # physical block sizes. Pick the common (smallest) block size so that + # KV-cache registration and transfer work correctly for both models. + backends = get_current_attn_backends(self.vllm_config) + kernel_block_size = select_common_block_size(self.block_size, backends) + if self.block_size != kernel_block_size: + logger.info_once( + "User-specified logical block size (%s) does not match" + " physical kernel block size (%s). Using the latter.", + self.block_size, + kernel_block_size, + ) + assert self.block_size > kernel_block_size + self.block_size = kernel_block_size + def __del__(self): self.shutdown() @@ -1268,9 +1288,6 @@ class MooncakeConnectorWorker: self.block_len_per_layer.append( curr_tensor_size_bytes // self.num_blocks ) - - kernel_block_size = cache.shape[-2 if self.use_mla else -3] - assert self.block_size == kernel_block_size kv_data_ptrs.append(base_addr) kv_data_lens.append(curr_tensor_size_bytes) From 102d51c9f3dc9696fdb88d36d301a65774e6ce59 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Wed, 15 Apr 2026 12:01:13 -0700 Subject: [PATCH 026/150] [CI] Only build release Docker images when NIGHTLY=1 (#39882) Signed-off-by: Kevin H. Luu Co-authored-by: Claude Opus 4.6 (1M context) --- .buildkite/release-pipeline.yaml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 45b2996f7ea..b3a6bb8ed4c 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -98,8 +98,15 @@ steps: commands: - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh" + - block: "Unblock to build release Docker images" + depends_on: ~ + key: block-build-release-images + if: build.env("NIGHTLY") != "1" + - group: "Build release Docker images" key: "build-release-images" + depends_on: block-build-release-images + allow_dependency_failure: true steps: - label: "Build release image - x86_64 - CUDA 12.9" depends_on: ~ @@ -617,6 +624,8 @@ steps: - label: ":docker: Build release image - x86_64 - ROCm" id: build-rocm-release-image depends_on: + - step: block-build-release-images + allow_failure: true - step: build-rocm-base-wheels allow_failure: false agents: From 41488f2acdc53eadbae98a316df47ba039589fe7 Mon Sep 17 00:00:00 2001 From: zhanqiuhu <49648934+ZhanqiuHu@users.noreply.github.com> Date: Wed, 15 Apr 2026 16:08:58 -0400 Subject: [PATCH 027/150] [Bugfix][NIXL] Fix `_logical_to_kernel_block_ids` conversion for non-mamba models (#39724) Signed-off-by: Zhanqiu Hu --- .../unit/test_nixl_connector_hma.py | 93 +++++++++++++++++++ .../kv_connector/v1/nixl/worker.py | 4 + 2 files changed, 97 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 17edd2069b3..0dfe4df7357 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -89,6 +89,99 @@ def test_logical_to_kernel_block_ids_with_hma(): ) +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "has_mamba,swa_enabled,mamba_enabled,remote_ratio," + "remote_block_ids,expected_remote_block_ids", + [ + # Non-mamba (FA+SWA): both groups expanded via _logical_to_kernel_block_ids. + # Regression for https://github.com/vllm-project/vllm/pull/39724 + ( + False, + True, + False, + 1, + ([0, 1, 2], [3, 4]), + [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]], + ), + # Mamba (FA+Mamba): FA expanded via _logical_to_remote_kernel_block_ids, + # Mamba passed through unchanged. + # remote_ratio=261 (Nemotron 30B TP=1) != local_ratio=2 so that using + # the wrong conversion method produces different FA results. + ( + True, + False, + True, + 261, + ([0, 1, 2], [10, 11]), + [[0, 1, 261, 262, 522, 523], [10, 11]], + ), + ], + ids=["non_mamba_fa_swa", "mamba_fa_ssm"], +) +def test_read_blocks_for_req_expands_remote_ids( + has_mamba, + swa_enabled, + mamba_enabled, + remote_ratio, + remote_block_ids, + expected_remote_block_ids, +): + """_read_blocks_for_req must expand remote logical block IDs to kernel + block IDs when kernel block size != logical block size. + + Non-mamba path uses _logical_to_kernel_block_ids (all groups expanded). + Mamba path uses _logical_to_remote_kernel_block_ids (FA expanded, Mamba + passed through). + """ + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + + worker = object.__new__(NixlConnectorWorker) + worker._has_mamba = has_mamba + worker._physical_blocks_per_logical_kv_block = 2 + worker.kv_cache_config = make_kv_cache_config( + block_size=16, swa_enabled=swa_enabled, mamba_enabled=mamba_enabled + ) + + remote_engine_id = "remote-engine" + if has_mamba: + worker._mamba_phys_ratio = {remote_engine_id: remote_ratio} + + # Mock kv_topo: empty remote ranks skips the transfer machinery entirely, + # isolating the block-ID expansion logic. + worker.kv_topo = MagicMock() + worker.kv_topo.get_target_remote_ranks_from_engine_id.return_value = [] + worker.kv_topo.tp_ratio_from_engine_id.return_value = 1 + + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="test-req", + local_block_ids=([0, 1], [2, 3]), + kv_transfer_params={ + "remote_block_ids": remote_block_ids, + "remote_engine_id": remote_engine_id, + "remote_request_id": "prefill-test-req", + "remote_host": "localhost", + "remote_port": 1234, + "tp_size": 1, + }, + ) + + meta = metadata.reqs_to_recv["test-req"] + worker._read_blocks_for_req("test-req", meta) + + assert meta.remote.block_ids == expected_remote_block_ids, ( + f"Expected {expected_remote_block_ids}, got {meta.remote.block_ids}" + ) + + @pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)]) def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size): """Test that a prefill instance returns fewer "remote blocks" for the SWA groups diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py index 4f47a8bc91a..45aa33033e7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py @@ -1914,6 +1914,10 @@ class NixlConnectorWorker: meta.remote.block_ids, self._mamba_phys_ratio[meta.remote.engine_id], ) + else: + meta.remote.block_ids = self._logical_to_kernel_block_ids( + meta.remote.block_ids + ) # D may have to perform multiple reads from different remote ranks. for i, remote_rank in enumerate(remote_ranks): if self.use_mla and tp_ratio < 0 and i > 0: From 0b790a25013e6b63a51ba00fc7da70537b3b3191 Mon Sep 17 00:00:00 2001 From: zhanqiuhu <49648934+ZhanqiuHu@users.noreply.github.com> Date: Wed, 15 Apr 2026 16:22:15 -0400 Subject: [PATCH 028/150] [Speculative Decoding] Add DFlash speculators config parsing (#38300) Signed-off-by: Zhanqiu Hu --- .buildkite/test_areas/spec_decode.yaml | 13 ++ .../v1/spec_decode/test_speculators_dflash.py | 171 ++++++++++++++++++ vllm/model_executor/models/qwen3_dflash.py | 9 +- .../configs/speculators/algos.py | 31 ++++ 4 files changed, 223 insertions(+), 1 deletion(-) create mode 100644 tests/v1/spec_decode/test_speculators_dflash.py diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index a0b73096867..76cc887ed0a 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -42,3 +42,16 @@ steps: - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" + +- label: DFlash Speculators Correctness + timeout_in_minutes: 30 + device: h100 + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/model_executor/models/qwen3_dflash.py + - tests/v1/spec_decode/test_speculators_dflash.py + commands: + - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 + - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test diff --git a/tests/v1/spec_decode/test_speculators_dflash.py b/tests/v1/spec_decode/test_speculators_dflash.py new file mode 100644 index 00000000000..2ba580695dd --- /dev/null +++ b/tests/v1/spec_decode/test_speculators_dflash.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline +from tests.utils import large_gpu_mark +from vllm import LLM +from vllm.config import SpeculativeConfig +from vllm.distributed import cleanup_dist_env_and_memory + +MODEL_PATH = "nm-testing/dflash-qwen3-8b-speculators" + +EXPECTED_GSM8K_ACCURACY = 0.885 +ACCURACY_RTOL = 0.03 +EXPECTED_ACCEPTANCE_LEN = 3.45 +ACCEPTANCE_LEN_RTOL = 0.15 + +# Expected per-position acceptance rates (accepted_at_pos / num_drafts) +# Based on GSM8K evaluation with Qwen3-8B dflash speculators. +EXPECTED_PER_POS_ACCEPTANCE_RATES = [0.795, 0.611, 0.429, 0.282] +PER_POS_RTOL = 0.15 + + +def compute_spec_decode_stats( + metrics, +) -> dict: + """Extract all spec-decode metrics and compute derived stats.""" + name2metric = {m.name: m for m in metrics} + + n_drafts = name2metric["vllm:spec_decode_num_drafts"].value + n_draft_tokens = name2metric["vllm:spec_decode_num_draft_tokens"].value + n_accepted = name2metric["vllm:spec_decode_num_accepted_tokens"].value + + per_pos_vec = name2metric["vllm:spec_decode_num_accepted_tokens_per_pos"].values + + acceptance_len = 1 + (n_accepted / n_drafts) if n_drafts > 0 else 1.0 + draft_tokens_per_step = (n_draft_tokens / n_drafts) if n_drafts > 0 else 0 + overall_acceptance_rate = (n_accepted / n_draft_tokens) if n_draft_tokens > 0 else 0 + per_pos_rates = [v / n_drafts for v in per_pos_vec] if n_drafts > 0 else [] + + return { + "num_drafts": n_drafts, + "num_draft_tokens": n_draft_tokens, + "num_accepted_tokens": n_accepted, + "acceptance_len": acceptance_len, + "draft_tokens_per_step": draft_tokens_per_step, + "overall_acceptance_rate": overall_acceptance_rate, + "per_pos_accepted": list(per_pos_vec), + "per_pos_acceptance_rates": per_pos_rates, + } + + +def print_spec_decode_stats(stats: dict) -> None: + """Print all spec-decode metrics and derived values.""" + print("\n===== Spec Decode Metrics =====") + print(f" num_drafts: {stats['num_drafts']}") + print(f" num_draft_tokens: {stats['num_draft_tokens']}") + print(f" num_accepted_tokens: {stats['num_accepted_tokens']}") + print(f" draft_tokens_per_step: {stats['draft_tokens_per_step']:.2f}") + print(f" overall_acceptance_rate: {stats['overall_acceptance_rate']:.4f}") + print(f" acceptance_len (1+acc/drafts): {stats['acceptance_len']:.4f}") + print(" per-position accepted tokens:", stats["per_pos_accepted"]) + print(" per-position acceptance rates:") + for i, rate in enumerate(stats["per_pos_acceptance_rates"]): + print(f" pos {i}: {rate:.4f}") + print("===============================\n") + + +def test_dflash_speculators_model(vllm_runner, example_prompts, monkeypatch): + """ + Test DFlash speculators model properly initializes speculative decoding. + + Verifies: + 1. Speculative config is automatically initialized from speculators config + 2. Method is detected as 'dflash' + 3. The draft model path is correctly set + 4. Speculative tokens count is valid (num_speculative_tokens=8) + 5. Text generation works with speculative decoding enabled + """ + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + + with vllm_runner( + MODEL_PATH, + dtype=torch.bfloat16, + enforce_eager=True, + quantization="fp8", + ) as vllm_model: + vllm_config = vllm_model.llm.llm_engine.vllm_config + + assert isinstance(vllm_config.speculative_config, SpeculativeConfig), ( + "Speculative config should be initialized for speculators model" + ) + + spec_config = vllm_config.speculative_config + assert spec_config.method == "dflash", ( + f"Expected method='dflash', got '{spec_config.method}'" + ) + assert spec_config.num_speculative_tokens > 0, ( + f"Expected positive speculative tokens, " + f"got {spec_config.num_speculative_tokens}" + ) + assert spec_config.model == MODEL_PATH, ( + f"Draft model should be {MODEL_PATH}, got {spec_config.model}" + ) + + vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens=20) + assert vllm_outputs, f"No outputs generated for speculators model {MODEL_PATH}" + + +@pytest.mark.slow_test +@large_gpu_mark(min_gb=40) +def test_dflash_speculators_correctness(monkeypatch): + """ + E2E correctness test for DFlash via the speculators auto-detect path. + + Evaluates GSM8k accuracy to ensure the speculators-format model produces + correct outputs, and checks that acceptance length does not collapse under + batched inference (lm-eval style). + + Observed per-position acceptance rates on GSM8K (1319 prompts): + pos 0: 0.795, pos 1: 0.611, pos 2: 0.429, pos 3: 0.282, + pos 4: 0.169, pos 5: 0.093, pos 6: 0.048, pos 7: 0.023 + Observed mean AL: 3.45 (GSM8K dataset, max_num_seqs=128) + """ + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + + spec_llm = LLM( + model=MODEL_PATH, + trust_remote_code=True, + max_model_len=4096, + max_num_seqs=128, + gpu_memory_utilization=0.85, + enforce_eager=False, + disable_log_stats=False, + ) + + results = evaluate_gsm8k_offline(spec_llm) + accuracy = results["accuracy"] + accuracy_threshold = EXPECTED_GSM8K_ACCURACY * (1 - ACCURACY_RTOL) + assert accuracy >= accuracy_threshold, ( + f"Expected GSM8K accuracy >= {accuracy_threshold:.3f}, got {accuracy:.3f}" + ) + + current_metrics = spec_llm.get_metrics() + stats = compute_spec_decode_stats(current_metrics) + print_spec_decode_stats(stats) + + acceptance_len = stats["acceptance_len"] + al_threshold = EXPECTED_ACCEPTANCE_LEN * (1 - ACCEPTANCE_LEN_RTOL) + assert acceptance_len >= al_threshold, ( + f"DFlash speculators acceptance length too low: " + f"{acceptance_len:.2f} < {al_threshold:.2f}" + ) + + # Check per-position acceptance rates for the first few positions. + per_pos_rates = stats["per_pos_acceptance_rates"] + for i, expected_rate in enumerate(EXPECTED_PER_POS_ACCEPTANCE_RATES): + assert i < len(per_pos_rates), ( + f"Missing per-position acceptance rate for position {i}" + ) + threshold = expected_rate * (1 - PER_POS_RTOL) + assert per_pos_rates[i] >= threshold, ( + f"Per-position acceptance rate at pos {i} too low: " + f"{per_pos_rates[i]:.4f} < {threshold:.4f} " + f"(expected ~{expected_rate:.4f})" + ) + + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index ce45136d7c0..cffe6267a4b 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -523,7 +523,14 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM): self.logits_processor = LogitsProcessor( self.config.draft_vocab_size, scale=logit_scale ) - self.draft_id_to_target_id = None + target_vocab_size = vllm_config.model_config.get_vocab_size() + if self.config.draft_vocab_size != target_vocab_size: + self.draft_id_to_target_id = nn.Parameter( + torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + requires_grad=False, + ) + else: + self.draft_id_to_target_id = None def embed_input_ids( self, diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py index f4dffab8b3e..405d5f5de1d 100644 --- a/vllm/transformers_utils/configs/speculators/algos.py +++ b/vllm/transformers_utils/configs/speculators/algos.py @@ -41,3 +41,34 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None: pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[ "eagle_aux_hidden_state_layer_ids" ] + + +@register_speculator("dflash") +def update_dflash(config_dict: dict, pre_trained_config: dict) -> None: + """ + Apply DFlash specific configuration transformations to the `dict` used to + construct the Transformers PreTrainedConfig. + + DFlash specific fields: + - draft_vocab_size: Size of the draft model's vocabulary + - target_hidden_size: Hidden size of the target model + - mask_token_id (required): Token ID used for parallel drafting mask + placeholders + - aux_hidden_state_layer_ids (required): Layer indices from the target + model whose intermediate hidden states are used as context for the + DFlash drafter. Mapped to both eagle_aux_hidden_state_layer_ids + (for gpu_model_runner) and dflash_config.target_layer_ids (for the + DFlash model). + """ + pre_trained_config["architectures"] = ["DFlashDraftModel"] + pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size") + if config_dict.get("target_hidden_size") is not None: + pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"] + + aux_layer_ids = config_dict["aux_hidden_state_layer_ids"] + pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids + + pre_trained_config["dflash_config"] = { + "mask_token_id": config_dict["mask_token_id"], + "target_layer_ids": aux_layer_ids, + } From 39ac640490ee2e8f951d343ae1707dd9bdacaf70 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Wed, 15 Apr 2026 16:28:43 -0400 Subject: [PATCH 029/150] [Bug] Fix batch invariant test issue, bs=1 with `max_seq_num = 1` (#39320) Signed-off-by: yewentao256 --- tests/v1/determinism/test_batch_invariance.py | 36 ++++++------------- 1 file changed, 11 insertions(+), 25 deletions(-) diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index b162b469bcd..eb5dec3e215 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -36,10 +36,10 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( using the high-level v1 LLM() API only (no manual batching). Strategy: - - Create two LLM engines with identical config except max_num_seqs: 1 vs N. - - Compute a baseline output for the needle prompt with the bs=1 engine. - - For many trials, generate a batch (size N) where the needle appears at a - random position among random filler prompts using the bs=N engine. + - Create a single LLM engine configured for the larger batch limit (N). + - Compute a baseline output for the needle prompt when it is run alone. + - For many trials, generate a mixed batch (size N) where the needle appears + at a random position among random filler prompts using the same engine. - Track how many trials match vs mismatch, and report totals at the end. The test fails if any mismatches occur, but we still dump pass/fail counts. @@ -83,11 +83,9 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( needle_prompt = "There once was a " - llm_bs1 = None - llm_bsN = None + llm = None try: - # Engine with bs=1 behavior - llm_bs1 = LLM_with_max_seqs( + llm = LLM_with_max_seqs( model=model, max_num_seqs=max_batch_size, gpu_memory_utilization=gpu_mem_util, @@ -96,20 +94,11 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( ) # Baseline generation for the needle prompt alone. - baseline_out = llm_bs1.generate([needle_prompt], sampling) + baseline_out = llm.generate([needle_prompt], sampling) assert len(baseline_out) == 1 assert len(baseline_out[0].outputs) >= 1 baseline_text = baseline_out[0].outputs[0].text - # Engine with larger batch limit (e.g., 64) - llm_bsN = LLM_with_max_seqs( - model=model, - max_num_seqs=max_batch_size, - gpu_memory_utilization=gpu_mem_util, - max_model_len=max_model_len, - attention_config=attention_config, - ) - mismatches = 0 for trial in range(num_trials): @@ -124,8 +113,8 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( else: prompts.append(_random_prompt(min_random_prompt, max_random_prompt)) - # Generate with the larger-batch engine - outputs = llm_bsN.generate(prompts, sampling) + # Generate with the same engine but in a larger batch. + outputs = llm.generate(prompts, sampling) # Find the needle output by position needle_output = outputs[needle_pos] assert needle_output.prompt == needle_prompt @@ -151,12 +140,9 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( finally: # Ensure engines are shutdown to free GPU/VRAM across test sessions - if llm_bs1 is not None: + if llm is not None: with contextlib.suppress(Exception): - llm_bs1.shutdown() - if llm_bsN is not None: - with contextlib.suppress(Exception): - llm_bsN.shutdown() + llm.shutdown() @skip_unsupported From ac3dac545b28ea6cf847e0044859e58f33d4f8b9 Mon Sep 17 00:00:00 2001 From: Benjamin Chislett Date: Wed, 15 Apr 2026 16:39:32 -0400 Subject: [PATCH 030/150] [Bugfix][Perf] Indexer upcast WK to BF16 for fusion (#38928) Signed-off-by: Benjamin Chislett --- vllm/model_executor/models/deepseek_mtp.py | 25 +++-- vllm/model_executor/models/deepseek_v2.py | 123 ++++++++++++--------- 2 files changed, 84 insertions(+), 64 deletions(-) diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 126efb6f88e..a66ec7aa3e6 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -30,6 +30,7 @@ from .deepseek_v2 import ( DeepseekV2DecoderLayer, DeepseekV2MixtureOfExperts, DeepseekV2MoE, + _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, ) from .utils import maybe_prefix @@ -190,10 +191,6 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): ) # Set MoE hyperparameters self.set_moe_parameters() - self.is_fp4_ckpt = ( - self.quant_config is not None - and self.quant_config.get_name() == "modelopt_fp4" - ) def set_moe_parameters(self): self.expert_weights = [] @@ -248,13 +245,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1), ] - if self.is_fp4_ckpt: - # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj) - indexer_fused_mapping = [ - ("wk_weights_proj", "wk", 0), - ("wk_weights_proj", "weights_proj", 1), - ] - stacked_params_mapping.extend(indexer_fused_mapping) + # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj) + indexer_fused_mapping = [ + ("wk_weights_proj", "wk", 0), + ("wk_weights_proj", "weights_proj", 1), + ] + stacked_params_mapping.extend(indexer_fused_mapping) expert_params_mapping = SharedFusedMoE.make_expert_params_mapping( self, @@ -271,6 +267,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + _pending_wk_fp8: dict = {} # FP8 indexer wk dequant buffer for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue @@ -281,6 +278,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) ) name = self._rewrite_spec_layer_name(spec_layer, name) + + if _try_load_fp8_indexer_wk( + name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params + ): + continue + for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if weight_name not in name: diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 17ddd5edece..cd28fb0192f 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -66,6 +66,10 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + scaled_dequantize, +) from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.sparse_attn_indexer import ( SparseAttnIndexer, @@ -628,10 +632,6 @@ class Indexer(nn.Module): self.vllm_config = vllm_config self.config = config self.quant_config = quant_config - self.is_fp4_ckpt = ( - self.quant_config is not None - and self.quant_config.get_name() == "modelopt_fp4" - ) # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"] self.topk_tokens = config.index_topk self.n_head = config.index_n_heads # 64 @@ -646,36 +646,16 @@ class Indexer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.wq_b", ) - if self.is_fp4_ckpt: - # Fused wk + weights_proj: single GEMM producing [head_dim + n_head]. - # weights_proj does not get quantized, - # so we run both with quant_config=None - # wk may be upcasted from the default quant; - # experiments show fusion is always faster unless WK proj is in FP4, - # which is not the case for all known quants. - self.wk_weights_proj = MergedColumnParallelLinear( - hidden_size, - [self.head_dim, self.n_head], - bias=False, - quant_config=None, - disable_tp=True, - prefix=f"{prefix}.wk_weights_proj", - ) - else: - self.wk = ReplicatedLinear( - hidden_size, - self.head_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.wk", - ) - self.weights_proj = ReplicatedLinear( - hidden_size, - self.n_head, - bias=False, - quant_config=None, - prefix=f"{prefix}.weights_proj", - ) + # Fused wk + weights_proj: single GEMM producing [head_dim + n_head]. + # FP8 wk weights are upcasted to BF16 during loading to maintain fusion. + self.wk_weights_proj = MergedColumnParallelLinear( + hidden_size, + [self.head_dim, self.n_head], + bias=False, + quant_config=None, + disable_tp=True, + prefix=f"{prefix}.wk_weights_proj", + ) self.k_norm = LayerNorm(self.head_dim, eps=1e-6) self.softmax_scale = self.head_dim**-0.5 @@ -716,14 +696,10 @@ class Indexer(nn.Module): q_pe, q_nope = torch.split( q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) - if self.is_fp4_ckpt: - # Fused wk + weights_proj: one GEMM, then split - kw, _ = self.wk_weights_proj(hidden_states) - k = kw[:, : self.head_dim] - weights = kw[:, self.head_dim :] - else: - k, _ = self.wk(hidden_states) - weights, _ = self.weights_proj(hidden_states) + # Fused wk + weights_proj: one GEMM, then split + kw, _ = self.wk_weights_proj(hidden_states) + k = kw[:, : self.head_dim] + weights = kw[:, self.head_dim :] k = self.k_norm(k) k_pe, k_nope = torch.split( @@ -761,6 +737,46 @@ class Indexer(nn.Module): return self.indexer_op(hidden_states, q_fp8, k, weights) +def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params): + """ + We fuse the WK and weights_proj projections, but in some checkpoints WK is stored + in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16. + Upcasting to BF16 during loading enables the fusion. This function loads the FP8 WK + weights and scale, and when both are available, dequantizes to BF16 and stores into + the fused wk_weights_proj.weight parameter. + """ + if "indexer.wk." not in name or "wk_weights" in name: + return False # Weight is not an isolated WK weight for the indexer, ignore. + is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn + is_scale = "weight_scale_inv" in name + if not is_weight and not is_scale: + return False # WK is not in FP8 format, ignore. + # Buffer this tensor (weight or scale) until both have arrived. + layer_prefix = name.rsplit(".wk.", 1)[0] # e.g. "model.layers.0.self_attn.indexer" + entry = buf.setdefault(layer_prefix, {}) + entry["weight" if is_weight else "scale"] = tensor + if "weight" not in entry or "scale" not in entry: + return True # still waiting for the other param + + # We have both weight and scale: dequantize FP8 to BF16. + weight_fp8, scale_inv = entry["weight"], entry["scale"] + del buf[layer_prefix] + block_size = weight_fp8.shape[1] // scale_inv.shape[1] + weight_bf16 = scaled_dequantize( + weight_fp8, + scale_inv, + group_shape=GroupShape(block_size, block_size), + out_dtype=torch.bfloat16, + ) + + # Load the dequantized weight into shard 0 of the fused buffer. + fused_name = f"{layer_prefix}.wk_weights_proj.weight" + param = params_dict[fused_name] + param.weight_loader(param, weight_bf16, 0) + loaded_params.add(fused_name) + return True + + def _min_latency_fused_qkv_a_proj_impl( input_: torch.Tensor, weight: torch.Tensor, @@ -1344,10 +1360,6 @@ class DeepseekV2ForCausalLM( quant_config = vllm_config.quant_config self.config = config self.quant_config = quant_config - self.is_fp4_ckpt = ( - self.quant_config is not None - and self.quant_config.get_name() == "modelopt_fp4" - ) qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0) qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0) @@ -1473,13 +1485,13 @@ class DeepseekV2ForCausalLM( ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), ] - if self.is_fp4_ckpt: - # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj) - indexer_fused_mapping = [ - ("wk_weights_proj", "wk", 0), - ("wk_weights_proj", "weights_proj", 1), - ] - stacked_params_mapping.extend(indexer_fused_mapping) + # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj) + _pending_wk_fp8: dict = {} # When WK is in FP8, we dequant to BF16 for fusion + indexer_fused_mapping = [ + ("wk_weights_proj", "wk", 0), + ("wk_weights_proj", "weights_proj", 1), + ] + stacked_params_mapping.extend(indexer_fused_mapping) if self.use_mha: stacked_params_mapping.extend(mha_params_mapping) @@ -1516,6 +1528,11 @@ class DeepseekV2ForCausalLM( rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) ) + if _try_load_fp8_indexer_wk( + name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params + ): + continue + for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if weight_name not in name: From c77e596e2eb6627d9f126436733550f2e31f642d Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Wed, 15 Apr 2026 16:43:15 -0400 Subject: [PATCH 031/150] [FlashAttention] Don't overwrite `flash_attn_interface.py` when installing precompiled (#39932) Signed-off-by: Matthew Bonanni --- setup.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 1748781985f..b0cca73bb91 100644 --- a/setup.py +++ b/setup.py @@ -693,6 +693,12 @@ class precompiled_wheel_utils: flash_attn_regex = re.compile( r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py" ) + # __init__.py and flash_attn_interface.py are source-controlled + # in vllm and should not be overwritten (matches cmake exclusions) + flash_attn_files_to_skip = { + "vllm/vllm_flash_attn/__init__.py", + "vllm/vllm_flash_attn/flash_attn_interface.py", + } triton_kernels_regex = re.compile( r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py" ) @@ -705,7 +711,11 @@ class precompiled_wheel_utils: filter(lambda x: x.filename in files_to_copy, wheel.filelist) ) file_members += list( - filter(lambda x: flash_attn_regex.match(x.filename), wheel.filelist) + filter( + lambda x: flash_attn_regex.match(x.filename) + and x.filename not in flash_attn_files_to_skip, + wheel.filelist, + ) ) file_members += list( filter( From 7c636432c62e242d36bff056d5258474498d3a4e Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Wed, 15 Apr 2026 17:20:06 -0400 Subject: [PATCH 032/150] [CI Bug] fix flaky test (#39938) Signed-off-by: yewentao256 --- tests/v1/kv_connector/unit/test_nixl_connector_hma.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 0dfe4df7357..30913ff98ee 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -195,7 +195,7 @@ def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size): llm_kwargs = { "model": model_name, "enforce_eager": True, - "gpu_memory_utilization": 0.5, + "gpu_memory_utilization": 0.47, "kv_transfer_config": kv_transfer_config, "max_model_len": 2048, # NOTE: Make sure HMA is enabled From 27c0ca50a0853e6087cd0cf8fe5caccc50f471f6 Mon Sep 17 00:00:00 2001 From: Collin McCarthy Date: Wed, 15 Apr 2026 16:09:11 -0700 Subject: [PATCH 033/150] Update registry for Nemotron-v3 VL Nano/Super (#39747) Signed-off-by: Collin McCarthy --- tests/models/registry.py | 44 ++++++++++++++++++++++++++ vllm/config/speculative.py | 4 +++ vllm/model_executor/models/registry.py | 2 ++ 3 files changed, 50 insertions(+) diff --git a/tests/models/registry.py b/tests/models/registry.py index 92ebac01841..9c15decd8fd 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1038,6 +1038,50 @@ _MULTIMODAL_EXAMPLE_MODELS = { }, 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", + 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, + ), + # 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, + ), "OpenCUAForConditionalGeneration": _HfExamplesInfo( "xlangai/OpenCUA-7B", trust_remote_code=True ), diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 61f311c8429..bbe923f68f1 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -300,6 +300,10 @@ class SpeculativeConfig: {"n_predict": n_predict, "architectures": ["ErnieMTPModel"]} ) + if hf_config.architectures[0] == "NemotronH_Super_Omni_Reasoning_V3": + # Promote VLM's text_config so MTP detection below fires correctly + hf_config = hf_config.text_config + if ( hf_config.model_type in {"nemotron_h", "nemotron_h_puzzle"} and hasattr(hf_config, "num_nextn_predict_layers") diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index ef89e159feb..01a767b9f32 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -473,6 +473,8 @@ _MULTIMODAL_MODELS = { "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"), "Molmo2ForConditionalGeneration": ("molmo2", "Molmo2ForConditionalGeneration"), "NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), + "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), + "NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), "NVLM_D": ("nvlm_d", "NVLM_D_Model"), "OpenCUAForConditionalGeneration": ("opencua", "OpenCUAForConditionalGeneration"), "OpenPanguVLForConditionalGeneration": ( From 6dc9491406938463e44e631610ce83df1d20f8cd Mon Sep 17 00:00:00 2001 From: Luciano Martins Date: Wed, 15 Apr 2026 20:13:07 -0300 Subject: [PATCH 034/150] [Model] Fix Gemma 4 token repetition by dynamic BOS injection for PT models (#39842) Signed-off-by: Luciano Martins Co-authored-by: Luciano Martins --- vllm/model_executor/models/gemma4_mm.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index e22f23c5c8b..d88a5f1725b 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -167,10 +167,15 @@ class Gemma4ProcessingInfo(BaseProcessingInfo): Setting ``add_special_tokens=False`` here prevents the duplicate and ensures both ``llm.generate()`` and the chat/completions API behave - correctly. + correctly for IT models. For PT models (without chat template), we + keep the default (True) to ensure BOS is added for raw prompts. """ + tokenizer = self.ctx.get_tokenizer() + has_chat_template = getattr(tokenizer, "chat_template", None) is not None + params = super().get_default_tok_params() - params = params.with_kwargs(add_special_tokens=False) + if has_chat_template: + params = params.with_kwargs(add_special_tokens=False) return params def get_hf_processor(self, **kwargs: object) -> Gemma4Processor: From 03f8d3a548ce9769f9fd89cb4505e8b77649c943 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Thu, 16 Apr 2026 00:29:15 +0100 Subject: [PATCH 035/150] Update to transformers v5 (#30566) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: khluu Signed-off-by: Kevin H. Luu Signed-off-by: Cyrus Leung Co-authored-by: khluu Co-authored-by: Cyrus Leung Co-authored-by: jiang1.li --- .../scripts/hardware_ci/run-cpu-test.sh | 2 +- .buildkite/test_areas/models_basic.yaml | 16 +- docker/Dockerfile | 9 +- docker/Dockerfile.cpu | 6 + docker/Dockerfile.nightly_torch | 7 +- docker/Dockerfile.rocm | 7 +- .../installation/gpu.rocm.inc.md | 2 +- requirements/common.txt | 4 +- requirements/test/cuda.in | 6 +- requirements/test/cuda.txt | 20 +-- requirements/test/nightly-torch.txt | 4 +- requirements/test/rocm.in | 5 +- requirements/test/rocm.txt | 31 ++-- requirements/test/xpu.in | 1 - requirements/test/xpu.txt | 24 +-- tests/conftest.py | 9 ++ tests/lora/conftest.py | 6 + tests/lora/test_minicpmv_tp.py | 11 ++ tests/model_executor/test_weight_utils.py | 18 --- .../models/language/generation/test_common.py | 5 + .../language/pooling_mteb_test/test_baai.py | 5 +- .../language/pooling_mteb_test/test_gte.py | 3 +- .../language/pooling_mteb_test/test_jina.py | 4 + .../multimodal/generation/test_common.py | 44 +++++- .../generation/test_nemotron_parse.py | 4 + .../multimodal/generation/test_phi4siglip.py | 11 ++ .../multimodal/generation/test_voxtral.py | 4 + .../multimodal/generation/vlm_utils/core.py | 5 + .../multimodal/pooling/test_colqwen3.py | 5 + .../multimodal/pooling/test_intern_vit.py | 5 + .../pooling/test_jinavl_reranker.py | 5 + .../processing/test_musicflamingo.py | 7 + tests/models/registry.py | 139 ++++++++++++++++-- tests/models/utils.py | 11 +- .../test_step3p5_reasoning_parser.py | 4 +- tests/v1/e2e/spec_decode/test_spec_decode.py | 6 +- .../model_loader/gguf_loader.py | 12 ++ vllm/model_executor/models/gemma4_mm.py | 51 +++++-- vllm/model_executor/models/musicflamingo.py | 2 +- .../models/transformers/base.py | 5 + vllm/tokenizers/registry.py | 35 ++++- 41 files changed, 445 insertions(+), 115 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index db75ad3083b..27ec0068668 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -16,5 +16,5 @@ echo "--- :docker: Building Docker image" docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu . # Run the image, setting --shm-size=4g for tensor parallel. -docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \ +docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 --shm-size=4g "$IMAGE_NAME" \ timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}" diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 10b038d8b8a..ed782c061fa 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -4,7 +4,6 @@ depends_on: steps: - label: Basic Models Tests (Initialization) timeout_in_minutes: 45 - device: h200_18gb torch_nightly: true source_file_dependencies: - vllm/ @@ -73,3 +72,18 @@ steps: - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl # Whisper needs spawn method to avoid deadlock - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper + +- label: Transformers Backward Compatibility Models Test + working_dir: "/vllm-workspace/" + optional: true + soft_fail: true + commands: + - pip install transformers==4.57.5 + - pytest -v -s tests/models/test_initialization.py + - pytest -v -s tests/models/test_transformers.py + - pytest -v -s tests/models/multimodal/processing/ + - pytest -v -s tests/models/multimodal/test_mapping.py + - python3 examples/offline_inference/basic/chat.py + - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl + # Whisper needs spawn method to avoid deadlock + - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper diff --git a/docker/Dockerfile b/docker/Dockerfile index 12942b5c807..3081d7ef138 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -642,7 +642,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ else \ BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \ fi; \ - uv pip install --system accelerate hf_transfer modelscope \ + uv pip install --system accelerate modelscope \ "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}" # ============================================================ @@ -756,9 +756,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system -e tests/vllm_test_utils # enable fast downloads from hf (for testing) -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system hf_transfer -ENV HF_HUB_ENABLE_HF_TRANSFER 1 +ENV HF_XET_HIGH_PERFORMANCE 1 + +# increase timeout for hf downloads (for testing) +ENV HF_HUB_DOWNLOAD_TIMEOUT 60 # Copy in the v1 package for testing (it isn't distributed yet) COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 0600f7da82f..77b449625dd 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -197,6 +197,12 @@ ADD ./.buildkite/ ./.buildkite/ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install -e tests/vllm_test_utils +# enable fast downloads from hf (for testing) +ENV HF_XET_HIGH_PERFORMANCE 1 + +# increase timeout for hf downloads (for testing) +ENV HF_HUB_DOWNLOAD_TIMEOUT 60 + ######################### RELEASE IMAGE ######################### FROM base AS vllm-openai diff --git a/docker/Dockerfile.nightly_torch b/docker/Dockerfile.nightly_torch index 39e1cc18759..0733509a0eb 100644 --- a/docker/Dockerfile.nightly_torch +++ b/docker/Dockerfile.nightly_torch @@ -272,9 +272,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system -e tests/vllm_test_utils # enable fast downloads from hf (for testing) -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system hf_transfer -ENV HF_HUB_ENABLE_HF_TRANSFER 1 +ENV HF_XET_HIGH_PERFORMANCE 1 + +# increase timeout for hf downloads (for testing) +ENV HF_HUB_DOWNLOAD_TIMEOUT 60 RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system -r requirements/test/nightly-torch.txt diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 801847d4999..fa7a5846edc 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -365,9 +365,10 @@ RUN cd /vllm-workspace \ && python3 -m pip install pytest-shard # enable fast downloads from hf (for testing) -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system hf_transfer -ENV HF_HUB_ENABLE_HF_TRANSFER=1 +ENV HF_XET_HIGH_PERFORMANCE=1 + +# increase timeout for hf downloads (for testing) +ENV HF_HUB_DOWNLOAD_TIMEOUT 60 # install audio decode package `torchcodec` from source (required due to # ROCm and torch version mismatch) for tests with datasets package diff --git a/docs/getting_started/installation/gpu.rocm.inc.md b/docs/getting_started/installation/gpu.rocm.inc.md index 4ab01ee8c68..f8385997eea 100644 --- a/docs/getting_started/installation/gpu.rocm.inc.md +++ b/docs/getting_started/installation/gpu.rocm.inc.md @@ -240,7 +240,7 @@ uv pip install vllm==${VLLM_VERSION} \ # Install dependencies pip install --upgrade numba \ scipy \ - huggingface-hub[cli,hf_transfer] \ + huggingface-hub[cli] \ setuptools_scm pip install -r requirements/rocm.txt diff --git a/requirements/common.txt b/requirements/common.txt index b610fd67868..299ec734ff3 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -7,7 +7,7 @@ requests >= 2.26.0 tqdm blake3 py-cpuinfo -transformers >= 4.56.0, < 5 +transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0 tokenizers >= 0.21.1 # Required for fast incremental detokenization. protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994 fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint. @@ -37,7 +37,7 @@ pyyaml six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12 setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12 einops # Required for Qwen2-VL. -compressed-tensors == 0.14.0.1 # required for compressed-tensors +compressed-tensors == 0.15.0.1 # required for compressed-tensors depyf==0.20.0 # required for profiling and debugging with compilation config cloudpickle # allows pickling lambda functions in model_executor/models/registry.py watchfiles # required for http server to monitor the updates of TLS files diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 378ecf94222..5cf3a69e1fb 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -18,7 +18,7 @@ httpx librosa # required for audio tests vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test -peft>=0.15.0 # required for phi-4-mm test +peft>=0.18.1 # required for phi-4-mm test pqdm ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests resampy # required for audio tests @@ -39,8 +39,8 @@ opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.11 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==4.57.5 -tokenizers==0.22.0 +transformers==5.5.3 +tokenizers==0.22.2 schemathesis>=3.39.15 # Required for openai schema test. # quantization bitsandbytes==0.49.2 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 548ca9310ff..ed67685e6eb 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -4,7 +4,7 @@ absl-py==2.1.0 # via # rouge-score # tensorboard -accelerate==1.0.1 +accelerate==1.13.0 # via peft aenum==3.1.16 # via lightly @@ -248,7 +248,6 @@ filelock==3.16.1 # huggingface-hub # ray # torch - # transformers # virtualenv fiona==1.10.1 # via torchgeo @@ -331,7 +330,7 @@ h5py==3.13.0 # via terratorch harfile==0.3.0 # via schemathesis -hf-xet==1.1.7 +hf-xet==1.4.3 # via huggingface-hub hiredis==3.0.0 # via tensorizer @@ -345,9 +344,10 @@ httpx==0.27.2 # via # -r requirements/test/cuda.in # diffusers + # huggingface-hub # perceptron # schemathesis -huggingface-hub==0.36.2 +huggingface-hub==1.10.2 # via # accelerate # datasets @@ -756,7 +756,7 @@ pathvalidate==3.2.1 # via pytablewriter patsy==1.0.1 # via statsmodels -peft==0.16.0 +peft==0.18.1 # via -r requirements/test/cuda.in perceptron==0.1.4 # via -r requirements/test/cuda.in @@ -982,7 +982,7 @@ referencing==0.35.1 # via # jsonschema # jsonschema-specifications -regex==2024.9.11 +regex==2026.2.28 # via # diffusers # nltk @@ -1002,7 +1002,6 @@ requests==2.32.3 # google-api-core # google-cloud-storage # gpt-oss - # huggingface-hub # lightly # lm-eval # mistral-common @@ -1015,7 +1014,6 @@ requests==2.32.3 # starlette-testclient # tacoreader # tiktoken - # transformers # wandb resampy==0.4.3 # via -r requirements/test/cuda.in @@ -1216,7 +1214,7 @@ timm==1.0.17 # segmentation-models-pytorch # terratorch # torchgeo -tokenizers==0.22.0 +tokenizers==0.22.2 # via # -c requirements/common.txt # -r requirements/test/cuda.in @@ -1295,7 +1293,7 @@ tqdm==4.67.3 # tacoreader # terratorch # transformers -transformers==4.57.5 +transformers==5.5.3 # via # -c requirements/common.txt # -r requirements/test/cuda.in @@ -1317,7 +1315,9 @@ typepy==1.3.2 typer==0.15.2 # via # fastsafetensors + # huggingface-hub # perceptron + # transformers types-python-dateutil==2.9.0.20241206 # via arrow typeshed-client==2.8.2 diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index e0eb7e11411..420fb496a71 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -29,8 +29,8 @@ opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.11 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==4.57.5 -tokenizers==0.22.0 +transformers==5.5.3 +tokenizers==0.22.2 schemathesis>=3.39.15 # Required for openai schema test. # quantization bitsandbytes>=0.49.2 diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index b5a9451b36f..dbb1500edcf 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -38,8 +38,8 @@ opencv-python-headless>=4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.11 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==4.57.5 -tokenizers==0.22.0 +transformers==5.5.3 +tokenizers==0.22.2 schemathesis>=3.39.15 # Required for openai schema test # quantization bitsandbytes==0.49.2 @@ -82,4 +82,3 @@ plotly # required for perf comparison html report rapidfuzz torchgeo==0.7.0 multiprocess==0.70.16 -huggingface-hub==0.36.2 diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index e1efae912ee..ba9cd3dfdcf 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -39,7 +39,7 @@ annotated-doc==0.0.4 # typer annotated-types==0.7.0 # via pydantic -anthropic==0.89.0 +anthropic==0.93.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -172,7 +172,7 @@ colorful==0.5.8 # via ray colorlog==6.10.1 # via optuna -compressed-tensors==0.14.0.1 +compressed-tensors==0.15.0.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -269,9 +269,9 @@ fastapi==0.135.2 # model-hosting-container-standards fastapi-cli==0.0.24 # via fastapi -fastapi-cloud-cli==0.15.1 +fastapi-cloud-cli==0.16.1 # via fastapi-cli -fastar==0.9.0 +fastar==0.10.0 # via fastapi-cloud-cli fastparquet==2026.3.0 # via genai-perf @@ -290,7 +290,6 @@ filelock==3.25.2 # python-discovery # ray # torch - # transformers # virtualenv fiona==1.10.1 # via torchgeo @@ -384,7 +383,7 @@ h5py==3.16.0 # via terratorch harfile==0.4.0 # via schemathesis -hf-xet==1.4.2 +hf-xet==1.4.3 # via huggingface-hub hiredis==3.3.1 # via tensorizer @@ -403,6 +402,7 @@ httpx==0.27.2 # diffusers # fastapi # fastapi-cloud-cli + # huggingface-hub # mcp # model-hosting-container-standards # openai @@ -410,9 +410,8 @@ httpx==0.27.2 # schemathesis httpx-sse==0.4.3 # via mcp -huggingface-hub==0.36.2 +huggingface-hub==1.10.2 # via - # -r requirements/test/rocm.in # accelerate # datasets # diffusers @@ -484,7 +483,7 @@ jinja2==3.1.6 # genai-perf # lm-eval # torch -jiter==0.13.0 +jiter==0.14.0 # via # anthropic # openai @@ -631,7 +630,7 @@ msgpack==1.1.2 # via # librosa # ray -msgspec==0.20.0 +msgspec==0.21.0 # via -r requirements/test/../common.txt mteb==2.11.5 # via -r requirements/test/rocm.in @@ -742,7 +741,7 @@ omegaconf==2.3.0 # lightning open-clip-torch==2.32.0 # via -r requirements/test/rocm.in -openai==2.30.0 +openai==2.31.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1093,7 +1092,7 @@ python-dotenv==1.2.2 # uvicorn python-json-logger==4.1.0 # via -r requirements/test/../common.txt -python-multipart==0.0.22 +python-multipart==0.0.26 # via # fastapi # mcp @@ -1180,7 +1179,6 @@ requests==2.32.5 # google-api-core # google-cloud-storage # gpt-oss - # huggingface-hub # lightly # lm-eval # mistral-common @@ -1194,7 +1192,6 @@ requests==2.32.5 # starlette-testclient # tacoreader # tiktoken - # transformers # wandb resampy==0.4.3 # via -r requirements/test/rocm.in @@ -1428,7 +1425,7 @@ timm==1.0.17 # segmentation-models-pytorch # terratorch # torchgeo -tokenizers==0.22.0 +tokenizers==0.22.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1471,7 +1468,7 @@ tqdm==4.67.3 # tacoreader # terratorch # transformers -transformers==4.57.5 +transformers==5.5.3 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1498,7 +1495,9 @@ typer==0.24.1 # fastapi-cli # fastapi-cloud-cli # fastsafetensors + # huggingface-hub # perceptron + # transformers typeshed-client==2.9.0 # via jsonargparse typing-extensions==4.15.0 diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index 0e4ca1d99dc..94ffc249395 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -13,7 +13,6 @@ pytest-shard absl-py accelerate arctic-inference -hf_transfer lm_eval[api] modelscope diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 51810592c46..4ddc0aa1c92 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -19,7 +19,9 @@ aiosignal==1.4.0 albumentations==1.4.6 # via -r requirements/test/xpu.in annotated-doc==0.0.4 - # via fastapi + # via + # fastapi + # typer annotated-types==0.7.0 # via pydantic anyio==4.13.0 @@ -64,6 +66,7 @@ click==8.3.1 # jiwer # nltk # schemathesis + # typer # uvicorn colorama==0.4.6 # via sacrebleu @@ -112,7 +115,6 @@ filelock==3.25.2 # huggingface-hub # modelscope # torch - # transformers frozenlist==1.8.0 # via # aiohttp @@ -133,9 +135,7 @@ h11==0.16.0 # uvicorn harfile==0.4.0 # via schemathesis -hf-transfer==0.1.9 - # via -r requirements/test/xpu.in -hf-xet==1.4.2 +hf-xet==1.4.3 # via huggingface-hub html2text==2025.4.15 # via gpt-oss @@ -144,8 +144,9 @@ httpcore==1.0.9 httpx==0.28.1 # via # datasets + # huggingface-hub # schemathesis -huggingface-hub==0.36.2 +huggingface-hub==1.10.2 # via # accelerate # datasets @@ -515,7 +516,6 @@ requests==2.33.1 # docker # evaluate # gpt-oss - # huggingface-hub # lm-eval # mistral-common # modelscope @@ -524,11 +524,11 @@ requests==2.33.1 # schemathesis # starlette-testclient # tiktoken - # transformers rich==14.3.3 # via # mteb # schemathesis + # typer rouge-score==0.1.2 # via lm-eval rpds-py==0.30.0 @@ -572,6 +572,8 @@ setuptools==80.10.2 # modelscope # pytablewriter # torch +shellingham==1.5.4 + # via typer six==1.17.0 # via # -c requirements/common.txt @@ -665,7 +667,7 @@ tqdm==4.67.3 # pqdm # sentence-transformers # transformers -transformers==4.57.6 +transformers==5.5.3 # via # -c requirements/common.txt # sentence-transformers @@ -676,6 +678,10 @@ typepy==1.3.4 # dataproperty # pytablewriter # tabledata +typer==0.24.1 + # via + # huggingface-hub + # transformers typing-extensions==4.15.0 # via # -c requirements/common.txt diff --git a/tests/conftest.py b/tests/conftest.py index a666c5a8663..bc657ff1ca7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -410,6 +410,15 @@ class HfRunner: model_name, trust_remote_code=trust_remote_code, ) + # HF runner should use the HF config so that it's consistent with the HF model + if self.config.__module__.startswith("vllm.transformers_utils.configs"): + from transformers.models.auto.configuration_auto import CONFIG_MAPPING + + del CONFIG_MAPPING._extra_content[self.config.model_type] + self.config = AutoConfig.from_pretrained( + model_name, + trust_remote_code=trust_remote_code, + ) self.device = self.get_default_device() self.dtype = dtype = _get_and_verify_dtype( self.model_name, diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py index 20944a9111e..169ddbf7ce5 100644 --- a/tests/lora/conftest.py +++ b/tests/lora/conftest.py @@ -3,6 +3,7 @@ import tempfile from collections import OrderedDict +from importlib import reload from unittest.mock import MagicMock import pytest @@ -47,6 +48,11 @@ def cleanup_fixture(should_do_global_cleanup_after_test: bool): def maybe_enable_lora_dual_stream(monkeypatch: pytest.MonkeyPatch): if current_platform.is_cuda(): monkeypatch.setenv("VLLM_LORA_ENABLE_DUAL_STREAM", "1") + import vllm.lora.layers.base_linear + + if not hasattr(vllm.lora.layers.base_linear, "lora_linear_async"): + # Reload the module to ensure the environment variable takes effect. + reload(vllm.lora.layers.base_linear) yield diff --git a/tests/lora/test_minicpmv_tp.py b/tests/lora/test_minicpmv_tp.py index e430826461a..3d6484a710a 100644 --- a/tests/lora/test_minicpmv_tp.py +++ b/tests/lora/test_minicpmv_tp.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from importlib.metadata import version + import pytest +from packaging.version import Version import vllm from vllm.assets.image import ImageAsset @@ -10,6 +13,14 @@ from vllm.platforms import current_platform from ..utils import multi_gpu_test +pytestmark = pytest.mark.skipif( + Version("5.0") <= Version(version("transformers")), + reason=( + "MiniCPMV custom processor uses tokenizer.im_start_id which is not " + "available on TokenizersBackend in transformers v5.0+" + ), +) + MODEL_PATH = "openbmb/MiniCPM-Llama3-V-2_5" PROMPT_TEMPLATE = ( diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py index 93535ae0aac..260ebdcefb3 100644 --- a/tests/model_executor/test_weight_utils.py +++ b/tests/model_executor/test_weight_utils.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os import tempfile import huggingface_hub.constants @@ -10,26 +9,10 @@ from huggingface_hub.utils import LocalEntryNotFoundError from vllm.model_executor.model_loader.weight_utils import ( download_weights_from_hf, - enable_hf_transfer, maybe_remap_kv_scale_name, ) -def test_hf_transfer_auto_activation(): - if "HF_HUB_ENABLE_HF_TRANSFER" in os.environ: - # in case it is already set, we can't test the auto activation - pytest.skip("HF_HUB_ENABLE_HF_TRANSFER is set, can't test auto activation") - enable_hf_transfer() - try: - # enable hf hub transfer if available - import hf_transfer # type: ignore # noqa - - HF_TRANSFER_ACTIVE = True - except ImportError: - HF_TRANSFER_ACTIVE = False - assert huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER == HF_TRANSFER_ACTIVE - - def test_download_weights_from_hf(): with tempfile.TemporaryDirectory() as tmpdir: # assert LocalEntryNotFoundError error is thrown @@ -178,5 +161,4 @@ class TestMaybeRemapKvScaleName: if __name__ == "__main__": - test_hf_transfer_auto_activation() test_download_weights_from_hf() diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py index c524480839b..b276f37a2a3 100644 --- a/tests/models/language/generation/test_common.py +++ b/tests/models/language/generation/test_common.py @@ -143,6 +143,11 @@ def test_models( # in parts of the operators pytest.skip(f"Skipping '{model}' model test with AITER kernel.") + if current_platform.is_cpu() and model == "TitanML/tiny-mixtral": + # This untrained model is sensitive to the rounding error + # Fuse ops to reduce bfloat16 rounding + monkeypatch.setenv("VLLM_CPU_CI_ENV", "0") + with hf_runner(model) as hf_model: hf_outputs = hf_model.generate_greedy_logprobs_limit( example_prompts, max_tokens, num_logprobs diff --git a/tests/models/language/pooling_mteb_test/test_baai.py b/tests/models/language/pooling_mteb_test/test_baai.py index 1199393d4b7..ec11960fda0 100644 --- a/tests/models/language/pooling_mteb_test/test_baai.py +++ b/tests/models/language/pooling_mteb_test/test_baai.py @@ -69,7 +69,10 @@ MODELS = [ attn_type="decoder", is_prefix_caching_supported=True, is_chunked_prefill_supported=True, - enable_test=True, + # Skip: model's custom tokenizer on HF hub is incompatible with + # transformers v5 (sets attrs before super().__init__, triggering + # AttributeError on 'verbose' in __getattr__). + enable_test=False, ), ] diff --git a/tests/models/language/pooling_mteb_test/test_gte.py b/tests/models/language/pooling_mteb_test/test_gte.py index 0c35d66c366..0a54262e124 100644 --- a/tests/models/language/pooling_mteb_test/test_gte.py +++ b/tests/models/language/pooling_mteb_test/test_gte.py @@ -72,7 +72,8 @@ MODELS = [ attn_type="encoder_only", is_prefix_caching_supported=False, is_chunked_prefill_supported=False, - enable_test=True, + # Skip: numerical regression with transformers v5. + enable_test=False, ), ########## ModernBertModel EmbedModelInfo( diff --git a/tests/models/language/pooling_mteb_test/test_jina.py b/tests/models/language/pooling_mteb_test/test_jina.py index 627cc043194..d75ec2a2ace 100644 --- a/tests/models/language/pooling_mteb_test/test_jina.py +++ b/tests/models/language/pooling_mteb_test/test_jina.py @@ -75,6 +75,10 @@ def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: mteb_test_rerank_models(vllm_runner, model_info) +@pytest.mark.skip( + reason="jinaai/jina-embeddings-v3 custom XLMRobertaLoRA model on HF hub " + "is incompatible with transformers v5 (missing all_tied_weights_keys)" +) @pytest.mark.parametrize("model_info", EMBEDDING_MODELS) @pytest.mark.parametrize("dtype", ["half"]) @pytest.mark.parametrize("dimensions", [16, 32]) diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index bf5119cf44f..1147ccef35b 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -186,7 +186,14 @@ VLM_TEST_SETTINGS = { max_num_seqs=2, auto_cls=AutoModel, hf_output_post_proc=model_utils.ultravox_trunc_hf_output, - marks=[pytest.mark.core_model, pytest.mark.cpu_model], + marks=[ + pytest.mark.core_model, + pytest.mark.cpu_model, + # TODO: Remove skip once model has been upstreamed to Transformers + pytest.mark.skip( + reason="Custom model code is not compatible with Transformers v5" + ), + ], ), #### Transformers fallback to test ## To reduce test burden, we only test batching arbitrary image size @@ -397,14 +404,14 @@ VLM_TEST_SETTINGS = { "gemma4": VLMTestInfo( models=["google/gemma-4-E2B-it"], test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - prompt_formatter=lambda img_prompt: f"user\n{img_prompt}\nmodel\n", # noqa: E501 + prompt_formatter=lambda img_prompt: f"<|turn>user\n{img_prompt}\n<|turn>model\n", # noqa: E501 single_image_prompts=IMAGE_ASSETS.prompts( { - "stop_sign": "What's the content in the center of the image?", - "cherry_blossom": "What is the season?", + "stop_sign": "<|image|>What's the content in the center of the image?", # noqa: E501 + "cherry_blossom": "<|image|>What is the season?", } ), - multi_image_prompt="Describe the two images in detail.", + multi_image_prompt="<|image|><|image|>Describe the two images in detail.", # noqa: E501 max_model_len=4096, max_num_seqs=2, auto_cls=AutoModelForImageTextToText, @@ -533,6 +540,12 @@ VLM_TEST_SETTINGS = { max_model_len=4096, use_tokenizer_eos=True, patch_hf_runner=model_utils.internvl_patch_hf_runner, + # TODO: Remove skip once model has been upstreamed to Transformers + marks=[ + pytest.mark.skip( + reason="Custom model code tries to access data from meta-tensor" + ) + ], ), "intern_vl-video": VLMTestInfo( models=[ @@ -545,6 +558,12 @@ VLM_TEST_SETTINGS = { use_tokenizer_eos=True, patch_hf_runner=model_utils.internvl_patch_hf_runner, num_logprobs=10 if current_platform.is_rocm() else 5, + # TODO: Remove skip once model has been upstreamed to Transformers + marks=[ + pytest.mark.skip( + reason="Custom model code tries to access data from meta-tensor" + ) + ], ), "intern_vl-hf": VLMTestInfo( models=["OpenGVLab/InternVL3-1B-hf"], @@ -591,6 +610,8 @@ VLM_TEST_SETTINGS = { hf_model_kwargs={"device_map": "auto"}, patch_hf_runner=model_utils.isaac_patch_hf_runner, image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], + # TODO: Remove skip once model has been upstreamed to Transformers + marks=[pytest.mark.skip(reason="Custom model imports deleted object")], # noqa: E501 ), "kimi_vl": VLMTestInfo( models=["moonshotai/Kimi-VL-A3B-Instruct"], @@ -806,7 +827,12 @@ VLM_TEST_SETTINGS = { pytest.mark.skipif( Version(TRANSFORMERS_VERSION) == Version("4.57.3"), reason="This model is broken in Transformers v4.57.3", - ) + ), + pytest.mark.skipif( + Version(TRANSFORMERS_VERSION) >= Version("5.0.0"), + reason="Model's custom code uses ROPE_INIT_FUNCTIONS" + "['default'] which was removed in transformers v5", + ), ], ), "phi3v": VLMTestInfo( @@ -960,6 +986,12 @@ VLM_TEST_SETTINGS = { ) for inp in custom_inputs.different_patch_input_cases_internvl() ], + # TODO: Remove skip once model has been upstreamed to Transformers + marks=[ + pytest.mark.skip( + reason="Custom model code tries to access data from meta-tensor" + ) + ], ), "llava_onevision-multiple-images": VLMTestInfo( models=["llava-hf/llava-onevision-qwen2-0.5b-ov-hf"], diff --git a/tests/models/multimodal/generation/test_nemotron_parse.py b/tests/models/multimodal/generation/test_nemotron_parse.py index e224f31e6df..8159cc9a8da 100644 --- a/tests/models/multimodal/generation/test_nemotron_parse.py +++ b/tests/models/multimodal/generation/test_nemotron_parse.py @@ -103,6 +103,10 @@ def run_test( ) +@pytest.mark.skip( + reason="Model's custom MBart decoder has head count mismatch with " + "transformers v5's GQA-aware cross-attention (8 vs 16 heads)" +) @pytest.mark.parametrize("model", ["nvidia/NVIDIA-Nemotron-Parse-v1.1"]) @pytest.mark.parametrize("dtype", ["bfloat16"]) @pytest.mark.parametrize("num_logprobs", [5]) diff --git a/tests/models/multimodal/generation/test_phi4siglip.py b/tests/models/multimodal/generation/test_phi4siglip.py index e8f4ba82925..f80b16c341b 100644 --- a/tests/models/multimodal/generation/test_phi4siglip.py +++ b/tests/models/multimodal/generation/test_phi4siglip.py @@ -2,9 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence +from importlib.metadata import version import pytest import regex as re +from packaging.version import Version from transformers import AutoModelForCausalLM, AutoTokenizer from vllm.logprobs import SampleLogprobs @@ -19,6 +21,15 @@ from ....conftest import ( from ....utils import multi_gpu_test from ...utils import check_logprobs_close +pytestmark = pytest.mark.skipif( + Version("5.0") <= Version(version("transformers")), + reason=( + "vllm upgraded transformers above v5.4 where HF model custom code uses siglip2 " + "internals (filter_out_non_signature_kwargs) removed by " + "huggingface/transformers#43514" + ), +) + MODEL_ID = "microsoft/Phi-4-reasoning-vision-15B" HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts( diff --git a/tests/models/multimodal/generation/test_voxtral.py b/tests/models/multimodal/generation/test_voxtral.py index 590b549dcf5..82db1dc6812 100644 --- a/tests/models/multimodal/generation/test_voxtral.py +++ b/tests/models/multimodal/generation/test_voxtral.py @@ -149,6 +149,10 @@ def test_online_serving(vllm_runner, audio_assets: AudioTestAssets): ) +@pytest.mark.skip( + reason="VoxtralProcessor.apply_chat_template() in transformers v5 " + "doesn't resolve chat_template=None to the default template" +) def test_hf_reference(hf_runner, vllm_runner, audio_assets: AudioTestAssets): """Compare vLLM Mistral-format output against HF Transformers reference. diff --git a/tests/models/multimodal/generation/vlm_utils/core.py b/tests/models/multimodal/generation/vlm_utils/core.py index 3de4ca209a6..ae95f39586c 100644 --- a/tests/models/multimodal/generation/vlm_utils/core.py +++ b/tests/models/multimodal/generation/vlm_utils/core.py @@ -80,6 +80,11 @@ def run_test( if vllm_runner_kwargs: vllm_runner_kwargs_.update(vllm_runner_kwargs) + # Avoid passing limit_mm_per_prompt twice when vllm_runner_kwargs + # already contains it (e.g. gemma4 sets it via vllm_runner_kwargs). + if "limit_mm_per_prompt" in vllm_runner_kwargs_: + limit_mm_per_prompt = vllm_runner_kwargs_.pop("limit_mm_per_prompt") + with vllm_runner( model, max_model_len=max_model_len, diff --git a/tests/models/multimodal/pooling/test_colqwen3.py b/tests/models/multimodal/pooling/test_colqwen3.py index 2faac7fbfb6..9eefedc153c 100644 --- a/tests/models/multimodal/pooling/test_colqwen3.py +++ b/tests/models/multimodal/pooling/test_colqwen3.py @@ -22,6 +22,11 @@ from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam from ....conftest import VllmRunner +pytestmark = pytest.mark.skip( + reason="ColQwen3 model's weight tying is incompatible with " + "transformers v5 (missing all_tied_weights_keys)" +) + MODELS = [ "TomoroAI/tomoro-colqwen3-embed-4b", "OpenSearch-AI/Ops-Colqwen3-4B", diff --git a/tests/models/multimodal/pooling/test_intern_vit.py b/tests/models/multimodal/pooling/test_intern_vit.py index c3f7c81b78b..d7b67b8bdb6 100644 --- a/tests/models/multimodal/pooling/test_intern_vit.py +++ b/tests/models/multimodal/pooling/test_intern_vit.py @@ -12,6 +12,11 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from ....conftest import ImageTestAssets +pytestmark = pytest.mark.skip( + reason="InternVisionModel's custom code is incompatible with " + "transformers v5 (missing all_tied_weights_keys)" +) + # we use snapshot_download to prevent conflicts between # dynamic_module and trust_remote_code for hf_runner DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"] diff --git a/tests/models/multimodal/pooling/test_jinavl_reranker.py b/tests/models/multimodal/pooling/test_jinavl_reranker.py index 035ca62058a..18a02625ea4 100644 --- a/tests/models/multimodal/pooling/test_jinavl_reranker.py +++ b/tests/models/multimodal/pooling/test_jinavl_reranker.py @@ -15,6 +15,11 @@ from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam from ....conftest import HfRunner, VllmRunner +pytestmark = pytest.mark.skip( + reason="jinaai/jina-reranker-m0 custom code is incompatible with " + "transformers v5 (missing all_tied_weights_keys)" +) + MODELS = ["jinaai/jina-reranker-m0"] MM_PROCESSOR_KWARGS = { diff --git a/tests/models/multimodal/processing/test_musicflamingo.py b/tests/models/multimodal/processing/test_musicflamingo.py index 625e1ad8d29..ba14b776029 100644 --- a/tests/models/multimodal/processing/test_musicflamingo.py +++ b/tests/models/multimodal/processing/test_musicflamingo.py @@ -17,11 +17,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from importlib.metadata import version from unittest.mock import MagicMock import numpy as np import pytest import torch +from packaging.version import Version from transformers import PretrainedConfig from tests.models.registry import HF_EXAMPLE_MODELS @@ -122,6 +124,11 @@ def test_musicflamingo_dummy_text_uses_plain_audio_tokens(mock_ctx): assert builder.get_dummy_text({"audio": 2}) == "" +@pytest.mark.skipif( + Version(version("transformers")) >= Version("5.5"), + reason="transformers v5.5 added native MusicFlamingoForConditionalGeneration " + "with a different get_audio_features signature (requires input_ids)", +) def test_musicflamingo_audio_feature_pipeline_matches_hf_small_config(): from transformers.models.musicflamingo import ( modeling_musicflamingo as hf_musicflamingo_modeling, diff --git a/tests/models/registry.py b/tests/models/registry.py index 9c15decd8fd..956565e551b 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -335,7 +335,15 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "internlm/internlm2-chat-7b", trust_remote_code=True ), "InternLM2VEForCausalLM": _HfExamplesInfo( - "OpenGVLab/Mono-InternVL-2B", trust_remote_code=True + "OpenGVLab/Mono-InternVL-2B", + trust_remote_code=True, + max_transformers_version="4.57", + transformers_version_reason={ + "vllm": ( + "Custom config cannot be loaded with Transformers " + "v5 because `vision_config` is not always set" + ) + }, ), "InternLM3ForCausalLM": _HfExamplesInfo( "internlm/internlm3-8b-instruct", trust_remote_code=True @@ -475,6 +483,13 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "Plamo2ForCausalLM": _HfExamplesInfo( "pfnet/plamo-2-1b", trust_remote_code=True, + max_transformers_version="4.57", + transformers_version_reason={ + "hf": ( + "Custom model code uses `_tied_weight_keys: list[str]` but " + "Transformers v5 now expects `_tied_weight_keys: dict[str, str]`" + ) + }, ), "Plamo3ForCausalLM": _HfExamplesInfo( "pfnet/plamo-3-nict-2b-base", @@ -515,6 +530,13 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { trust_remote_code=True, max_model_len=4096, is_available_online=True, + max_transformers_version="5.3", + transformers_version_reason={ + "vllm": ( + "vllm upgraded transformers above v5.4 where " + "validate_rope() no longer accepts ignore_keys param" + ) + }, ), "SeedOssForCausalLM": _HfExamplesInfo( "ByteDance-Seed/Seed-OSS-36B-Instruct", @@ -553,6 +575,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "xverse/XVERSE-7B-Chat", tokenizer="meta-llama/Llama-2-7b", trust_remote_code=True, + max_transformers_version="4.57", + transformers_version_reason={ + "vllm": "XVERSE tokenizer is incompatible with transformers v5 " + "(add_prefix_space / prepend_scheme mismatch).", + }, ), "Zamba2ForCausalLM": _HfExamplesInfo("Zyphra/Zamba2-7B-instruct"), "MiMoForCausalLM": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL", trust_remote_code=True), @@ -763,10 +790,18 @@ _MULTIMODAL_EXAMPLE_MODELS = { # [Decoder-only] "AriaForConditionalGeneration": _HfExamplesInfo("rhymes-ai/Aria"), "AudioFlamingo3ForConditionalGeneration": _HfExamplesInfo( - "nvidia/audio-flamingo-3-hf", min_transformers_version="5.0.0" + "nvidia/audio-flamingo-3-hf", + min_transformers_version="5.3.0", + transformers_version_reason={ + "vllm": "Needs https://github.com/huggingface/transformers/pull/43538" + }, ), "MusicFlamingoForConditionalGeneration": _HfExamplesInfo( - "nvidia/music-flamingo-2601-hf", min_transformers_version="5.3.0" + "nvidia/music-flamingo-2601-hf", + min_transformers_version="5.3.0", + transformers_version_reason={ + "vllm": "Needs https://github.com/huggingface/transformers/pull/43538" + }, ), "AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereLabs/aya-vision-8b"), "BagelForConditionalGeneration": _HfExamplesInfo("ByteDance-Seed/BAGEL-7B-MoT"), @@ -821,12 +856,30 @@ _MULTIMODAL_EXAMPLE_MODELS = { ), "FireRedASR2ForConditionalGeneration": _HfExamplesInfo( "allendou/FireRedASR2-LLM-vllm", + trust_remote_code=True, + max_transformers_version="5.1", + transformers_version_reason={ + "vllm": "Incompatible with transformers v5.2+ " + "(dict object has no attribute '__name__').", + }, ), "FireRedLIDForConditionalGeneration": _HfExamplesInfo( "PatchyTisa/FireRedLID-vllm", + trust_remote_code=True, + max_transformers_version="5.1", + transformers_version_reason={ + "vllm": "Incompatible with transformers v5.2+ " + "(dict object has no attribute '__name__').", + }, ), "FunASRForConditionalGeneration": _HfExamplesInfo( "allendou/Fun-ASR-Nano-2512-vllm", + trust_remote_code=True, + max_transformers_version="5.1", + transformers_version_reason={ + "vllm": "Incompatible with transformers v5.2+ " + "(dict object has no attribute '__name__').", + }, ), "FunAudioChatForConditionalGeneration": _HfExamplesInfo( "funaudiochat", is_available_online=False @@ -868,6 +921,13 @@ _MULTIMODAL_EXAMPLE_MODELS = { "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", @@ -887,7 +947,12 @@ _MULTIMODAL_EXAMPLE_MODELS = { extras={"0.2-2B-Preview": "PerceptronAI/Isaac-0.2-2B-Preview"}, ), "InternS1ForConditionalGeneration": _HfExamplesInfo( - "internlm/Intern-S1", trust_remote_code=True + "internlm/Intern-S1", + trust_remote_code=True, + max_transformers_version="4.57", + transformers_version_reason={ + "vllm": "Custom tokenizer code is not compatible with Transformers v5." + }, ), "InternS1ProForConditionalGeneration": _HfExamplesInfo( "internlm/Intern-S1-Pro", @@ -976,7 +1041,14 @@ _MULTIMODAL_EXAMPLE_MODELS = { "MiDashengLMModel": _HfExamplesInfo( "mispeech/midashenglm-7b", trust_remote_code=True ), - "MiniCPMO": _HfExamplesInfo("openbmb/MiniCPM-o-2_6", trust_remote_code=True), + "MiniCPMO": _HfExamplesInfo( + "openbmb/MiniCPM-o-2_6", + trust_remote_code=True, + max_transformers_version="4.57", + transformers_version_reason={ + "hf": "Custom processor code is not compatible with Transformers v5." + }, + ), "MiniCPMV": _HfExamplesInfo( "openbmb/MiniCPM-Llama3-V-2_5", extras={ @@ -984,6 +1056,13 @@ _MULTIMODAL_EXAMPLE_MODELS = { "4.0": "openbmb/MiniCPM-V-4", "4.5": "openbmb/MiniCPM-V-4_5", }, + max_transformers_version="4.57", + transformers_version_reason={ + "vllm": ( + "MiniCPMVBatchFeature is incompatible with its base class in " + "Transformers v5. See https://huggingface.co/openbmb/MiniCPM-Llama3-V-2_5/discussions/78" + ) + }, trust_remote_code=True, ), "MiniMaxVL01ForConditionalGeneration": _HfExamplesInfo( @@ -1083,13 +1162,25 @@ _MULTIMODAL_EXAMPLE_MODELS = { trust_remote_code=True, ), "OpenCUAForConditionalGeneration": _HfExamplesInfo( - "xlangai/OpenCUA-7B", trust_remote_code=True + "xlangai/OpenCUA-7B", + trust_remote_code=True, + max_transformers_version="4.57", + transformers_version_reason={ + "vllm": "Tokenizer cannot be initialised in Transformers v5." + }, ), "OpenPanguVLForConditionalGeneration": _HfExamplesInfo( "FreedomIntelligence/openPangu-VL-7B", trust_remote_code=True, max_model_len=4096, enforce_eager=True, + max_transformers_version="4.57", + transformers_version_reason={ + "vllm": ( + "OpenPanguVLVideoProcessorInitKwargs does not specify total=False, " + "making all kwargs required. See https://huggingface.co/FreedomIntelligence/openPangu-VL-7B/discussions/2" + ) + }, ), "Ovis": _HfExamplesInfo( "AIDC-AI/Ovis2-1B", @@ -1101,12 +1192,24 @@ _MULTIMODAL_EXAMPLE_MODELS = { "1.6-gemma": "AIDC-AI/Ovis1.6-Gemma2-9B", }, ), - "Ovis2_5": _HfExamplesInfo("AIDC-AI/Ovis2.5-2B", trust_remote_code=True), + "Ovis2_5": _HfExamplesInfo( + "AIDC-AI/Ovis2.5-2B", + trust_remote_code=True, + max_transformers_version="4.57", + transformers_version_reason={ + "vllm": "Custom processor code is not compatible with Transformers v5." + }, + ), "Ovis2_6ForCausalLM": _HfExamplesInfo( "AIDC-AI/Ovis2.6-2B", is_available_online=False, trust_remote_code=True ), "Ovis2_6_MoeForCausalLM": _HfExamplesInfo( - "AIDC-AI/Ovis2.6-30B-A3B", trust_remote_code=True + "AIDC-AI/Ovis2.6-30B-A3B", + trust_remote_code=True, + max_transformers_version="4.57", + transformers_version_reason={ + "vllm": "Custom processor code is not compatible with Transformers v5." + }, ), "PaddleOCRVLForConditionalGeneration": _HfExamplesInfo( "PaddlePaddle/PaddleOCR-VL", @@ -1126,7 +1229,17 @@ _MULTIMODAL_EXAMPLE_MODELS = { extras={"phi3.5": "microsoft/Phi-3.5-vision-instruct"}, ), "Phi4ForCausalLMV": _HfExamplesInfo( - "microsoft/Phi-4-reasoning-vision-15B", trust_remote_code=True + "microsoft/Phi-4-reasoning-vision-15B", + trust_remote_code=True, + max_transformers_version="5.3", + transformers_version_reason={ + "vllm": ( + "vllm upgraded transformers above v5.4 where HF model " + "custom code uses siglip2 internals " + "(filter_out_non_signature_kwargs) removed " + "by huggingface/transformers#43514" + ) + }, ), "Phi4MMForCausalLM": _HfExamplesInfo( "microsoft/Phi-4-multimodal-instruct", trust_remote_code=True @@ -1223,6 +1336,14 @@ _MULTIMODAL_EXAMPLE_MODELS = { "architectures": ["Tarsier2ForConditionalGeneration"], "model_type": "tarsier2", }, + max_transformers_version="5.3", + transformers_version_reason={ + "vllm": ( + "Qwen2VLConfig was split into Qwen2VLConfig + " + "Qwen2VLTextConfig in transformers v5, breaking " + "attribute access (num_attention_heads, hidden_size, etc.)" + ) + }, ), "VoxtralForConditionalGeneration": _HfExamplesInfo( "mistralai/Voxtral-Mini-3B-2507", diff --git a/tests/models/utils.py b/tests/models/utils.py index 3b94f34fab0..b93beee6aa3 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -476,7 +476,16 @@ def dummy_hf_overrides( else: # Use minimal layers for testing num_layers = 1 - num_hidden_layers = 3 if model_arch == "Gemma3nForConditionalGeneration" else 1 + num_hidden_layers = ( + 3 + if model_arch + in ( + "Gemma3nForConditionalGeneration", + "Gemma4ForCausalLM", + "Gemma4ForConditionalGeneration", + ) + else 1 + ) update_dict = { "num_layers": num_layers, diff --git a/tests/reasoning/test_step3p5_reasoning_parser.py b/tests/reasoning/test_step3p5_reasoning_parser.py index 2196d247cb4..8f62e7a2cb4 100644 --- a/tests/reasoning/test_step3p5_reasoning_parser.py +++ b/tests/reasoning/test_step3p5_reasoning_parser.py @@ -2,10 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest -from transformers import AutoTokenizer from tests.reasoning.utils import run_reasoning_extraction from vllm.reasoning import ReasoningParser, ReasoningParserManager +from vllm.tokenizers import get_tokenizer parser_name = "step3p5" start_token = "" @@ -16,7 +16,7 @@ REASONING_MODEL_NAME = "stepfun-ai/Step-3.5-Flash" @pytest.fixture(scope="module") def step3p5_tokenizer(): - return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME) + return get_tokenizer(tokenizer_name=REASONING_MODEL_NAME) SIMPLE_REASONING = { diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py index c11bdbc50f7..a8fed766528 100644 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ b/tests/v1/e2e/spec_decode/test_spec_decode.py @@ -557,12 +557,16 @@ def test_eagle_correctness_light( "auto", 0.8, ), - ( + pytest.param( ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), False, False, "transformers", 0.8, + # TODO(hmellor): figure out why memory usage is so high + marks=pytest.mark.skip( + reason="Feature is experimental and uses too much memory in CI", + ), ), pytest.param( ( diff --git a/vllm/model_executor/model_loader/gguf_loader.py b/vllm/model_executor/model_loader/gguf_loader.py index ce6a813b8da..fc6f88b49ee 100644 --- a/vllm/model_executor/model_loader/gguf_loader.py +++ b/vllm/model_executor/model_loader/gguf_loader.py @@ -265,12 +265,24 @@ class GGUFModelLoader(BaseModelLoader): GGUF tensor name with suffix (e.g., 'mm.soft_emb_norm.weight') or None if no mapping found """ + # In transformers v5, multimodal models (e.g. Gemma3) wrap + # all sub-models under an outer 'model.' attribute, producing + # state_dict keys like 'model.language_model.layers.0...' and + # 'model.vision_tower.vision_model...'. Strip this outer + # prefix so the keys match what gguf-py expects. + if is_multimodal and hf_name.startswith("model."): + hf_name = hf_name[6:] # Remove outer 'model.' + # Strip 'language_model.' prefix for multimodal models - gguf-py # tensor mappings expect parameter names without this prefix. # Note: 'model.' prefix should be KEPT for text-only models as # gguf-py expects it. if hf_name.startswith("language_model."): hf_name = hf_name[15:] # Remove 'language_model.' + # Re-add 'model.' prefix because gguf-py text tensor maps + # expect 'model.layers...' format. + if is_multimodal: + hf_name = "model." + hf_name # Parse parameter name and suffix if hf_name.endswith((".weight", ".bias")): diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index d88a5f1725b..cd0ca37616b 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -125,8 +125,12 @@ class Gemma4AudioInputs(TensorSchema): """ type: Literal["audio"] = "audio" - input_features_padded: Annotated[torch.Tensor, TensorShape("bn", "s", "f")] - input_features_mask: Annotated[torch.Tensor, TensorShape("bn", "s")] + input_features_padded: Annotated[ + torch.Tensor, TensorShape("bn", "s", "f", dynamic_dims={"s"}) + ] + input_features_mask: Annotated[ + torch.Tensor, TensorShape("bn", "s", dynamic_dims={"s"}) + ] Gemma4ImageInputs = Gemma4ImagePixelInputs @@ -510,6 +514,8 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]): video_timestamps_per_video: list[list[float]] = [] video_frame_counts: list[int] = [] + video_replacements: list[str] = [] + for item in videos: video_array, metadata = item @@ -562,10 +568,7 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]): video_timestamps_per_video.append(timestamps) video_frame_counts.append(len(frames)) - # Build expanded replacement text and replace the - # <|video|> placeholder in the prompt. - # Use split(token, 1) to avoid collision — the - # replacement text itself contains <|video|> tokens. + # Build expanded replacement text for this video. ts_strs = [f"{int(s // 60):02d}:{int(s % 60):02d}" for s in timestamps] replacement = " ".join( f"{t} {processor.boi_token}" @@ -573,9 +576,23 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]): f"{processor.eoi_token}" for t, n in zip(ts_strs, num_soft_per_frame) ) - parts = prompt.split(processor.video_token, 1) - if len(parts) == 2: - prompt = parts[0] + replacement + parts[1] + video_replacements.append(replacement) + + # Replace all <|video|> placeholders at once. We split on + # video_token to get N+1 parts, then interleave with the + # N replacement strings. This avoids the iterative + # split-replace bug where replacement text (which itself + # contains <|video|> tokens) collides with later splits. + vt = processor.video_token + parts = prompt.split(vt, len(video_replacements)) + + # NOTE: len(parts) <= len(video_replacements) + 1 + parts_with_repl: list[str] = [] + for part, repl in zip(parts, video_replacements): + parts_with_repl.extend([part, repl]) + parts_with_repl.extend(parts[len(video_replacements) :]) + + prompt = "".join(parts_with_repl) video_outputs = { "pixel_values_videos": torch.cat(all_video_pixel_values, dim=0), @@ -638,19 +655,23 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]): ) if "input_features" in processed_outputs: - # Keep padded features for batched audio tower execution. - processed_outputs["input_features_padded"] = processed_outputs[ - "input_features" - ] - # Unpad per-item so each item's cache entry is self-contained. + # Unpad per-item so each item's cache entry is + # self-contained. The batched() field config in + # _get_mm_fields_config will re-pad all fields to the + # batch's max length at batch time, ensuring consistent + # padding regardless of cache history. + masks = processed_outputs["input_features_mask"] unpadded_features = [ f[mask] for f, mask in zip( processed_outputs["input_features"], - processed_outputs["input_features_mask"], + masks, ) ] + unpadded_masks = [mask[mask] for mask in masks] processed_outputs["input_features"] = unpadded_features + processed_outputs["input_features_padded"] = unpadded_features + processed_outputs["input_features_mask"] = unpadded_masks # Merge video outputs into the final result combined_outputs = dict(processed_outputs, **video_outputs) diff --git a/vllm/model_executor/models/musicflamingo.py b/vllm/model_executor/models/musicflamingo.py index f4e3bbe379a..497b2e63a7e 100644 --- a/vllm/model_executor/models/musicflamingo.py +++ b/vllm/model_executor/models/musicflamingo.py @@ -32,9 +32,9 @@ from transformers.models.musicflamingo import ( from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions +from vllm.inputs import MultiModalDataDict from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( - MultiModalDataDict, MultiModalFieldConfig, MultiModalKwargsItems, ) diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index b7967985f22..a3e4b844b80 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -275,6 +275,11 @@ class Base( ) class SupportTorchCompileWrapper(cls): ... + # Preserve __module__ so transformers v5's source-file checks + # (e.g. _can_set_experts_implementation) read the original + # model's module instead of this file. + SupportTorchCompileWrapper.__module__ = cls.__module__ + # Patch the class in its module module = sys.modules[cls.__module__] setattr(module, cls.__name__, SupportTorchCompileWrapper) diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index 7d48e3c6ff9..8f16e6d28f4 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib from dataclasses import dataclass, field from functools import lru_cache from pathlib import Path @@ -10,6 +11,7 @@ from typing_extensions import TypeVar, assert_never import vllm.envs as envs from vllm.logger import init_logger +from vllm.transformers_utils.config import get_config from vllm.transformers_utils.gguf_utils import ( check_gguf_file, get_gguf_file_path_from_hf, @@ -31,6 +33,13 @@ if TYPE_CHECKING: logger = init_logger(__name__) +# Model types whose hub tokenizer_class is incorrect and should be overridden with +# TokenizersBackend (the generic fast tokenizer). Adding a model type here is always a +# temporary workaround and better long term solutions are: +# - Add model type to MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS in transformers (better) +# - Fix tokenizer_class on the hub for the affected models (best) +_MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl"} + _VLLM_TOKENIZERS = { "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"), "grok2": ("grok2", "Grok2Tokenizer"), @@ -202,7 +211,31 @@ def get_tokenizer( **kwargs, ) - if tokenizer_cls == TokenizerLike: + # Ensure that, if the config were to come from vllm.transformers_utils.config, it is + # registered with AutoConfig before the tokenizer is loaded. This is necessary since + # tokenizer_cls_.from_pretrained will call AutoConfig.from_pretrained internally. + # This may fail for paths that don't have a model config (e.g. LoRA adapters), + # which is fine — those don't need custom config registration. + config = None + with contextlib.suppress(ValueError, OSError): + config = get_config( + tokenizer_name, + trust_remote_code=trust_remote_code, + revision=revision, + ) + + # Some models have an incorrect tokenizer_class on the hub. + # For these model types, bypass AutoTokenizer and use TokenizersBackend directly. + model_type = getattr(config, "model_type", None) if config else None + if model_type in _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: + from transformers.tokenization_utils_tokenizers import TokenizersBackend + + logger.debug( + "Overriding tokenizer_class to TokenizersBackend for model_type=%r", + model_type, + ) + tokenizer_cls_ = TokenizersBackend + elif tokenizer_cls == TokenizerLike: tokenizer_cls_ = TokenizerRegistry.load_tokenizer_cls(tokenizer_mode) else: tokenizer_cls_ = tokenizer_cls From 19fa90ed0d616f3c0bc0dbb3bf6bfb7772df7f32 Mon Sep 17 00:00:00 2001 From: Asaf Gardin <39553475+Josephasafg@users.noreply.github.com> Date: Thu, 16 Apr 2026 04:03:32 +0300 Subject: [PATCH 036/150] [Quantization] - Layerwise reloading of Attention/KV quantized models (#38995) Signed-off-by: Josephasafg --- .../model_loader/test_reload.py | 28 ++++++ .../model_loader/reload/__init__.py | 5 +- .../model_loader/reload/layerwise.py | 88 ++++++++++++++----- .../model_loader/weight_utils.py | 4 +- 4 files changed, 98 insertions(+), 27 deletions(-) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index fc79f6a0228..6e3e2d63e14 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -164,6 +164,34 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner): assert add_perp < mul_perp +def test_kv_scale_reload(vllm_runner): + """Test reloading a checkpoint that contains k_scale/v_scale weights.""" + if not current_platform.supports_fp8(): + pytest.skip(reason="Requires FP8 support") + + model = "nm-testing/Llama-3.2-1B-Instruct-FP8-KV" + + # Load dummy weights, then reload real checkpoint + with vllm_runner( + model_name=model, + load_format="dummy", + enable_prefix_caching=False, + max_model_len=16, + max_num_seqs=1, + ) as llm: + llm.collective_rpc( + "update_config", + kwargs={"overrides": {"load_config": {"load_format": "auto"}}}, + ) + llm.collective_rpc("reload_weights", kwargs={"weights_path": model}) + reloaded_perp = llm.generate_prompt_perplexity( + ["The capital of France is the city of Paris"], + mask=["The capital of France is"], + )[0] + + assert reloaded_perp < 10 + + @pytest.mark.parametrize( "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])] ) diff --git a/vllm/model_executor/model_loader/reload/__init__.py b/vllm/model_executor/model_loader/reload/__init__.py index 56a9d88ac4e..61dc1b66e6f 100644 --- a/vllm/model_executor/model_loader/reload/__init__.py +++ b/vllm/model_executor/model_loader/reload/__init__.py @@ -8,10 +8,9 @@ which is useful for weight updates without full model reconstruction. Limitations: 1. Composition with CPU offloading has not been implemented -2. Reloading Attention/MLA weights (q_scale, k_scale, v_scale) has not been implemented -3. Tied parameters will only reflect processing from one of the parent layers (for +2. Tied parameters will only reflect processing from one of the parent layers (for example, only processing from embed_tokens will have an effect) -4. This design assumes that the number of weights loaded from disk is the same as the +3. This design assumes that the number of weights loaded from disk is the same as the number of weights created at model init time. This is not true for quant methods which (1) pad weights or (2) load qkv weights into the same parameter. Both of these cases are non-issues for today's quant methods, but future quantizations may cause diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 2934b8b5ad5..2ebe2544478 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -200,6 +200,8 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon if hasattr(model, "_original_do_torchao_reload"): model._do_torchao_reload = model._original_do_torchao_reload + deferred_attn: list[tuple[torch.nn.Module, LayerReloadingInfo]] = [] + for layer in model.modules(): info = get_layerwise_info(layer) if not info.can_load(): @@ -208,22 +210,11 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon # Attention/MLA layers are processed after all other layers if isinstance(layer, (Attention, MLAAttention)): - if info.load_numel > 0: - raise NotImplementedError( - "Layerwise reloading of Q/K/V scale weights is not implemented yet" - ) - - elif info.kernel_tensors is None: - raise NotImplementedError( - "Layerwise loading of Q/K/V scale weights is not implemented yet" - ) - - else: - _place_kernel_tensors(layer, info) - layer.process_weights_after_loading(model_config.dtype) + deferred_attn.append((layer, info)) + continue # No weights were loaded - elif info.load_numel <= 0: + if info.load_numel <= 0: # first load: checkpoint did not contain weights for this layer if info.kernel_tensors is None: _layerwise_process(layer, info) @@ -244,11 +235,58 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon info.reset() + # Process attention layers after all other layers are done + for layer, info in deferred_attn: + _finalize_attention_layer(layer, info, model_config) + info.reset() + def finalize_layerwise_reload(*args, **kwargs): finalize_layerwise_processing(*args, **kwargs) +def _finalize_attention_layer( + layer: torch.nn.Module, info: LayerReloadingInfo, model_config: ModelConfig +) -> None: + if info.load_numel > 0 and info.kernel_tensors is not None: + # Reload with new scale weights from checkpoint + _place_kernel_tensors(layer, info) + _reload_attention_scales(layer, info) + elif info.load_numel > 0 or info.kernel_tensors is None: + raise ValueError( + "Layerwise loading of attention layers is not supported. " + "Attention must always process after linears." + ) + else: + _place_kernel_tensors(layer, info) + layer.process_weights_after_loading(model_config.dtype) + + +def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) -> None: + """Load and process attention scale weights (k_scale, v_scale, etc.) + during reload. + + Assumes dtype/shapes of attention tensors do not change during + processing, since we use .data.copy_() to preserve kernel tensor + references.""" + quant_method = getattr(layer, "quant_method", None) + if quant_method is None: + return + + # Re-create scale Parameters with sentinel values so unloaded scales + # are correctly detected by process_weights_after_loading + quant_method.create_weights(layer) + + for name, args in info.loaded_weights: + param = getattr(layer, name) + args.arguments["param"] = param + _get_weight_loader(param)(*args.args, **args.kwargs) + + quant_method.process_weights_after_loading(layer) + + _copy_and_restore_kernel_tensors(layer, info) + + def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): """ Finalize layer loading after all weights have been buffered. @@ -278,7 +316,6 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): param.weight_loader(*args.args, **args.kwargs) # Process weights (quantization, repacking, etc.) - # Attention/MLA are processed in `finalize_layerwise_reload` quant_method = getattr(layer, "quant_method", None) if isinstance(quant_method, QuantizeMethodBase): quant_method.process_weights_after_loading(layer) @@ -286,13 +323,7 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): # Copy processed values into original tensor storage (preserves cudagraph refs) # this code is a no-op if not reloading (because kernel tensors is empty) if info.kernel_tensors is not None: - parameters, buffers = info.kernel_tensors - for name, param in parameters.items(): - param.data.copy_(getattr(layer, name)) - for name, buffer in buffers.items(): - buffer.data.copy_(getattr(layer, name)) - - _place_kernel_tensors(layer, info) + _copy_and_restore_kernel_tensors(layer, info) info.reset() logger.debug("%s: Processed", layer.__class__.__name__) @@ -311,6 +342,19 @@ def _get_weight_loader(tensor: torch.Tensor): return getattr(tensor, "weight_loader", default_weight_loader) +def _copy_and_restore_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo): + """Copy processed values into original kernel tensor storage and restore + kernel tensor references on the layer. Preserves cudagraph references.""" + assert info.kernel_tensors is not None + parameters, buffers = info.kernel_tensors + for name, param in parameters.items(): + param.data.copy_(getattr(layer, name)) + for name, buffer in buffers.items(): + buffer.data.copy_(getattr(layer, name)) + + _place_kernel_tensors(layer, info) + + def _place_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo): for name in get_layer_tensors(layer): delattr(layer, name) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 3b961e8e143..8282e6b099d 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -1364,8 +1364,8 @@ def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> N if param.numel() == 1 and loaded_weight.numel() == 1: # Sometimes scalar values aren't considered tensors with shapes # so if both param and loaded_weight are a scalar, - # "broadcast" instead of copy - param.data.fill_(loaded_weight.item()) + # reshape to match before copying + param.data.copy_(loaded_weight.view(param.shape)) else: assert param.size() == loaded_weight.size(), ( f"Attempted to load weight ({loaded_weight.size()}) " From 343f65234bb2f6a0c42bf398e80a1a79b9739aaa Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Wed, 15 Apr 2026 18:09:11 -0700 Subject: [PATCH 037/150] =?UTF-8?q?[Model=20Runner=20V2][BugFix]=20fix=20n?= =?UTF-8?q?um=5Fsampled=20dtype=20for=20probabilistic=20rej=E2=80=A6=20(#3?= =?UTF-8?q?9951)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Giancarlo Delfin --- .../gpu/spec_decode/probabilistic_rejection_sampler_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 22f5d9df0d0..2feaae24539 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 @@ -527,7 +527,7 @@ def probabilistic_rejection_sample( sampled = draft_sampled.new_empty( num_reqs, num_speculative_steps + 1, dtype=torch.int64 ) - num_sampled = sampled.new_empty(num_reqs) + num_sampled = sampled.new_empty(num_reqs, dtype=torch.int32) target_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32) draft_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32) _probabilistic_rejection_kernel[(num_reqs,)]( From 5f7fab881a13df2926f50cb8d9f67b0bbc2ee41f Mon Sep 17 00:00:00 2001 From: vllmellm Date: Thu, 16 Apr 2026 09:55:29 +0800 Subject: [PATCH 038/150] [ROCm][FEAT] Integrate aiter gemm w8a8 ptpc (#33773) Signed-off-by: vllmellm --- tests/utils.py | 1 + vllm/_aiter_ops.py | 106 ++++++++++- .../model_executor/kernels/linear/__init__.py | 12 +- .../kernels/linear/scaled_mm/aiter.py | 164 +++++++++++++++++- .../schemes/compressed_tensors_w8a8_fp8.py | 9 +- .../layers/quantization/fbgemm_fp8.py | 2 + .../model_executor/layers/quantization/fp8.py | 4 +- .../layers/quantization/modelopt.py | 2 + .../quark/schemes/quark_w8a8_fp8.py | 2 + 9 files changed, 279 insertions(+), 23 deletions(-) diff --git a/tests/utils.py b/tests/utils.py index cfa5f525f5d..d35555d37c6 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1856,6 +1856,7 @@ class TestFP8Layer(torch.nn.Module): out_dtype=out_dtype, force_kernel=force_kernel, ) + self.kernel.process_weights_after_loading(self) def is_quant_fp8_enabled(self) -> bool: return self.kernel.quant_fp8.enabled() diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 8c2659b9c7e..d71e210fb51 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -3,6 +3,7 @@ import functools from collections.abc import Callable +import pandas as pd import torch from torch._ops import OpOverload @@ -56,6 +57,29 @@ def is_aiter_found_and_supported() -> bool: return False +@functools.cache +def _load_gemm_tuned_configs( + q_dtype_w: torch.dtype, csv_path: str +) -> set[tuple[int, int, int]]: + try: + df = pd.read_csv(csv_path).drop_duplicates() + df = df[df["q_dtype_w"] == str(q_dtype_w)] + return set(zip(df["N"].astype(int), df["K"].astype(int), df["M"].astype(int))) + except Exception: + return set() + + +def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool: + configs = _load_gemm_tuned_configs(q_dtype_w, csv_path) + l_m = ( + [1, 2, 4] + + list(range(8, 513, 8)) + + [1024, 1536] + + [2**i for i in range(11, 19)] + ) + return any((N, K, M) in configs for M in l_m) + + def if_aiter_supported(func: Callable) -> Callable: """Decorator that only executes the function if ROCm AITER package is supported and enabled on gfx9 archs. @@ -468,7 +492,7 @@ def _rocm_aiter_mla_decode_fwd_fake( pass -def _rocm_aiter_gemm_a8w8_impl( +def _rocm_aiter_w8a8_gemm_impl( A: torch.Tensor, B: torch.Tensor, As: torch.Tensor, @@ -485,7 +509,7 @@ def _rocm_aiter_gemm_a8w8_impl( return gemm_a8w8_CK(A, B, As, Bs, bias, output_dtype) -def _rocm_aiter_gemm_a8w8_fake( +def _rocm_aiter_w8a8_gemm_fake( A: torch.Tensor, B: torch.Tensor, As: torch.Tensor, @@ -499,6 +523,35 @@ def _rocm_aiter_gemm_a8w8_fake( return Y +def _rocm_aiter_preshuffled_per_token_w8a8_gemm_impl( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + from aiter import gemm_a8w8_bpreshuffle + + output = gemm_a8w8_bpreshuffle(A, B, As, Bs, None, output_dtype) + if bias is not None: + output.add_(bias) + return output + + +def _rocm_aiter_preshuffled_per_token_w8a8_gemm_fake( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + m = A.shape[0] + n = B.shape[0] + return torch.empty(m, n, dtype=output_dtype, device=A.device) + + def _rocm_aiter_triton_gemm_a8w8_blockscale_impl( A: torch.Tensor, B: torch.Tensor, @@ -1313,11 +1366,15 @@ class rocm_aiter_ops: ) direct_register_custom_op( - op_name="rocm_aiter_gemm_a8w8", - op_func=_rocm_aiter_gemm_a8w8_impl, - mutates_args=[], - fake_impl=_rocm_aiter_gemm_a8w8_fake, - dispatch_key=current_platform.dispatch_key, + op_name="rocm_aiter_w8a8_gemm", + op_func=_rocm_aiter_w8a8_gemm_impl, + fake_impl=_rocm_aiter_w8a8_gemm_fake, + ) + + direct_register_custom_op( + op_name="_rocm_aiter_preshuffled_per_token_w8a8_gemm", + op_func=_rocm_aiter_preshuffled_per_token_w8a8_gemm_impl, + fake_impl=_rocm_aiter_preshuffled_per_token_w8a8_gemm_fake, ) direct_register_custom_op( @@ -1493,7 +1550,7 @@ class rocm_aiter_ops: ) @staticmethod - def gemm_a8w8( + def w8a8_gemm( A: torch.Tensor, B: torch.Tensor, As: torch.Tensor, @@ -1501,7 +1558,20 @@ class rocm_aiter_ops: bias: torch.Tensor | None = None, output_dtype: torch.dtype = torch.float16, ) -> torch.Tensor: - return torch.ops.vllm.rocm_aiter_gemm_a8w8(A, B, As, Bs, bias, output_dtype) + return torch.ops.vllm.rocm_aiter_w8a8_gemm(A, B, As, Bs, bias, output_dtype) + + @staticmethod + def preshuffled_per_token_w8a8_gemm( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float16, + ) -> torch.Tensor: + return torch.ops.vllm._rocm_aiter_preshuffled_per_token_w8a8_gemm( + A, B, As, Bs, bias, output_dtype + ) @staticmethod def triton_gemm_a8w8_blockscale( @@ -1920,6 +1990,24 @@ class rocm_aiter_ops: (8192, 3584), ] + @staticmethod + def is_shuffled_per_token_w8a8_gemm_tuned( + N: int, K: int, q_dtype_w: torch.dtype + ) -> bool: + import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops + + csv_path = ( + aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE + ) + return _check_kernel_tuned(N, K, q_dtype_w, csv_path) + + @staticmethod + def is_per_token_w8a8_gemm_tuned(N: int, K: int, q_dtype_w: torch.dtype) -> bool: + import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops + + csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE + return _check_kernel_tuned(N, K, q_dtype_w, csv_path) + @staticmethod def shuffle_weight( tensor: torch.Tensor, layout: tuple[int, int] = (16, 16) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index e0a9272c890..6cbb65e26d6 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -106,6 +106,8 @@ from vllm.model_executor.kernels.linear.scaled_mm import ( from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( AiterFp8BlockScaledMMKernel, AiterInt8ScaledMMLinearKernel, + AiterPerTokenFp8ScaledMMLinearKernel, + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( CPUInt8ScaledMMLinearKernel, @@ -165,6 +167,8 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = ChannelWiseTorchFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, + AiterPerTokenFp8ScaledMMLinearKernel, ROCmFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, RowWiseTorchFP8ScaledMMLinearKernel, @@ -360,18 +364,18 @@ def choose_scaled_mm_linear_kernel( def init_fp8_linear_kernel( activation_quant_key: QuantKey, weight_quant_key: QuantKey, - weight_shape: tuple[int, int], input_dtype: torch.dtype, out_dtype: torch.dtype, - force_kernel: type[_KernelT] | None = None, + weight_shape: tuple[int, int], + force_kernel: type[FP8ScaledMMLinearKernel] | None = None, module_name: str | None = None, ) -> FP8ScaledMMLinearKernel | Fp8BlockScaledMMLinearKernel: scaled_mm_linear_kernel_config = FP8ScaledMMLinearLayerConfig( weight_quant_key=weight_quant_key, activation_quant_key=activation_quant_key, - weight_shape=weight_shape, input_dtype=input_dtype, out_dtype=out_dtype, + weight_shape=weight_shape, ) if activation_quant_key.scale.group_shape.is_per_group(): @@ -725,6 +729,8 @@ __all__ = [ "FP8ScaledMMLinearLayerConfig", "Int8ScaledMMLinearLayerConfig", "ScaledMMLinearLayerConfig", + "AiterPreshuffledPerTokenFp8ScaledMMLinearKernel", + "AiterPerTokenFp8ScaledMMLinearKernel", "NvFp4LinearKernel", "NvFp4LinearLayerConfig", "AiterInt8ScaledMMLinearKernel", diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 01d2298ed4a..8a8650d2213 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -5,18 +5,27 @@ import torch from vllm import _custom_ops as ops -from vllm._aiter_ops import rocm_aiter_ops +from vllm._aiter_ops import ( + rocm_aiter_ops, +) +from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, ) +from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from .BlockScaledMMLinearKernel import ( Fp8BlockScaledMMLinearKernel, - FP8ScaledMMLinearLayerConfig, ) from .cutlass import CutlassInt8ScaledMMLinearKernel -from .ScaledMMLinearKernel import Int8ScaledMMLinearLayerConfig +from .ScaledMMLinearKernel import ( + FP8ScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, + Int8ScaledMMLinearLayerConfig, +) + +logger = init_logger(__name__) class AiterInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel): @@ -113,7 +122,154 @@ class AiterInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel): # a to be [M, K] # b to be [N, K] # CutlassInt8ScaledMMLinearKernel prepare weight `w_q` in [K, N] format - return rocm_aiter_ops.gemm_a8w8(x_q, w_q.t(), x_s, w_s, bias, out_dtype) + return rocm_aiter_ops.w8a8_gemm(x_q, w_q.t(), x_s, w_s, bias, out_dtype) + + +class AiterPreshuffledPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_rocm(): + return False, "requires ROCm." + if not rocm_aiter_ops.is_linear_fp8_enabled(): + return ( + False, + "requires setting `VLLM_ROCM_USE_AITER=1` " + "and `VLLM_ROCM_USE_AITER_LINEAR=1`. " + "`VLLM_ROCM_USE_AITER_LINEAR` default is True.", + ) + try: + import aiter # noqa: F401 + except Exception: + return False, "requires aiter library to be installed." + return True, None + + @classmethod + def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: + is_ptpc = ( + c.activation_quant_key.scale.group_shape.is_per_token() + and c.weight_quant_key.scale.group_shape.is_per_channel() + ) + if c.weight_shape is None: + return False, "weight_shape is required for Aiter kernels" + N, K = c.weight_shape + fp8_dtype = current_platform.fp8_dtype() + + if c.out_dtype is not torch.bfloat16: + return False, "requires bfloat16 output dtype." + + if not is_ptpc: + return ( + False, + "requires per token activation scales and per channel weight scales.", + ) + + if not (N % 16 == 0 and K % 16 == 0): + return ( + False, + f"requires N and K dimensions divisible by 16, received " + f"N={N} and K={K}.", + ) + + # Aiter's shuffled per-token Gemm performs better than torch only when its + # tuned. + if not rocm_aiter_ops.is_shuffled_per_token_w8a8_gemm_tuned(N, K, fp8_dtype): + return ( + False, + f"requires a tuned configuration for N: {N} and K: {K} " + f"and fp8 dtype {fp8_dtype}.", + ) + + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + w_name, *_ = self.layer_param_names + w, *_ = self._get_layer_params(layer) + + replace_parameter( + layer, + w_name, + torch.nn.Parameter( + rocm_aiter_ops.shuffle_weight(w.t().contiguous()).data, + requires_grad=False, + ), + ) + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + return rocm_aiter_ops.preshuffled_per_token_w8a8_gemm( + A, B, As, Bs, bias, out_dtype + ) + + +class AiterPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported( + compute_capability + ) + + @classmethod + def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: + is_ptpc = ( + c.activation_quant_key.scale.group_shape.is_per_token() + and c.weight_quant_key.scale.group_shape.is_per_channel() + ) + if c.weight_shape is None: + return False, "weight_shape is required for Aiter kernels" + N, K = c.weight_shape + fp8_dtype = current_platform.fp8_dtype() + + if not is_ptpc: + return ( + False, + "requires per token activation scales and per channel weight scales.", + ) + + # Aiter's per-token Gemm performs better than torch oonly when its + # tuned. + if not rocm_aiter_ops.is_per_token_w8a8_gemm_tuned(N, K, fp8_dtype): + return ( + False, + f"requires a tuned configuration for N: {N} and K: {K} " + f"and fp8 dtype {fp8_dtype}.", + ) + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + w_name, *_ = self.layer_param_names + w, *_ = self._get_layer_params(layer) + + replace_parameter( + layer, + w_name, + torch.nn.Parameter(w.t(), requires_grad=False), + ) + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + return rocm_aiter_ops.w8a8_gemm(A, B, As, Bs, bias, out_dtype) class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py index 3bf606ddb33..7445634a825 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py @@ -31,8 +31,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, create_fp8_quant_key, kFp8DynamicTokenSym, + kFp8StaticChannelSym, kFp8StaticTensorSym, - kFp8StaticTokenSym, ) from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( cutlass_block_fp8_supported, @@ -47,7 +47,7 @@ activation_quant_key_mapping = { DYNAMIC_QUANT: kFp8DynamicTokenSym, } weight_quant_key_mapping = { - QuantizationStrategy.CHANNEL: kFp8StaticTokenSym, + QuantizationStrategy.CHANNEL: kFp8StaticChannelSym, QuantizationStrategy.TENSOR: kFp8StaticTensorSym, } logger = init_logger(__name__) @@ -67,7 +67,6 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enabled() assert not self.is_static_input_scheme self.act_q_group_shape = GroupShape(1, self.weight_block_size[0]) - self.weight_quant_key = create_fp8_quant_key( static=True, group_shape=GroupShape(*self.weight_block_size) ) @@ -76,7 +75,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): ) else: self.activation_quant_key = activation_quant_key_mapping[ - is_static_input_scheme + self.is_static_input_scheme ] self.weight_quant_key = weight_quant_key_mapping[self.strategy] @@ -138,9 +137,9 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): self.fp8_linear = init_fp8_linear_kernel( activation_quant_key=self.activation_quant_key, weight_quant_key=self.weight_quant_key, - weight_shape=layer.weight.shape, input_dtype=self.input_dtype, out_dtype=self.out_dtype, + weight_shape=(output_size_per_partition, input_size_per_partition), module_name=self.__class__.__name__, ) diff --git a/vllm/model_executor/layers/quantization/fbgemm_fp8.py b/vllm/model_executor/layers/quantization/fbgemm_fp8.py index 377cec364bf..d95c51be010 100644 --- a/vllm/model_executor/layers/quantization/fbgemm_fp8.py +++ b/vllm/model_executor/layers/quantization/fbgemm_fp8.py @@ -175,6 +175,8 @@ class FBGEMMFp8LinearMethod(LinearMethodBase): # Activations not quantized for marlin. del layer.input_scale_ub + self.fp8_linear.process_weights_after_loading(layer) + def apply( self, layer: torch.nn.Module, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index d7920462e61..94ed4f7c97b 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -397,8 +397,6 @@ class Fp8LinearMethod(LinearMethodBase): if self.block_quant: assert not self.act_q_static - self.fp8_linear.process_weights_after_loading(layer) - # If checkpoint not serialized fp8, quantize the weights. else: # If checkpoint is fp8 per-tensor, handle that there are N scales for N @@ -428,6 +426,8 @@ class Fp8LinearMethod(LinearMethodBase): else: layer.input_scale = None + self.fp8_linear.process_weights_after_loading(layer) + def apply( self, layer: torch.nn.Module, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 0b8ad0cbc1e..852ed1a10a3 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -517,6 +517,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase): layer.weight = Parameter(weight.t(), requires_grad=False) layer.weight_scale = Parameter(max_w_scale, requires_grad=False) layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) + self.fp8_linear.process_weights_after_loading(layer) def apply( self, @@ -597,6 +598,7 @@ class ModelOptFp8PcPtLinearMethod(LinearMethodBase): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight = Parameter(layer.weight.t(), requires_grad=False) layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False) + self.fp8_linear.process_weights_after_loading(layer) def apply( self, diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py index 3312e6901d6..6d94e26f960 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py @@ -120,6 +120,8 @@ class QuarkW8A8Fp8(QuarkScheme): if self.is_static_input_scheme: layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) + self.fp8_linear.process_weights_after_loading(layer) + def create_weights( self, layer: torch.nn.Module, From 951dca8019f6cf55f0fd210a004d3f3edfa0d58b Mon Sep 17 00:00:00 2001 From: Zhengxu Chen Date: Thu, 16 Apr 2026 00:03:41 -0400 Subject: [PATCH 039/150] [compile] Invoke split FX graph by codegen. (#38657) Signed-off-by: zhxchen17 --- vllm/compilation/backends.py | 25 +++++- vllm/compilation/caching.py | 35 ++++++-- vllm/compilation/codegen.py | 155 +++++++++++++++++++++++++++++++++++ 3 files changed, 206 insertions(+), 9 deletions(-) create mode 100644 vllm/compilation/codegen.py diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 7373ad75117..c3900ffc67d 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -1263,6 +1263,23 @@ class VllmBackend: original_split_gm if envs.VLLM_USE_MEGA_AOT_ARTIFACT else self.graph ) + from vllm.compilation.codegen import ( + compile_execution_fn, + generate_execution_code, + ) + + execution_code, submod_names = generate_execution_code(self.split_gm) + # Use getattr to get correct callables: __dict__ has PiecewiseBackend + # instances (from PiecewiseCompileInterpreter), _modules has originals. + # getattr checks __dict__ first, then falls back to _modules. + submod_callables = { + name: getattr(self.split_gm, name) + for name, _ in self.split_gm.named_children() + } + runtime_callable = compile_execution_fn( + execution_code, submod_callables, submod_names + ) + if ( self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE or not self.compilation_config.cudagraph_copy_inputs @@ -1271,9 +1288,11 @@ class VllmBackend: graph_to_serialize, example_inputs, self.prefix, - self.split_gm, + runtime_callable, is_encoder=self.is_encoder, vllm_backend=self, + execution_code=execution_code, + submod_names=submod_names, ) # index of tensors that have symbolic shapes (batch size) @@ -1294,7 +1313,7 @@ class VllmBackend: copy_and_call = make_copy_and_call( sym_tensor_indices, [example_inputs[x].clone() for x in sym_tensor_indices], - self.split_gm, + runtime_callable, ) return VllmSerializableFunction( @@ -1305,4 +1324,6 @@ class VllmBackend: is_encoder=self.is_encoder, vllm_backend=self, sym_tensor_indices=sym_tensor_indices, + execution_code=execution_code, + submod_names=submod_names, ) diff --git a/vllm/compilation/caching.py b/vllm/compilation/caching.py index c089f02a37f..6b61c0c770b 100644 --- a/vllm/compilation/caching.py +++ b/vllm/compilation/caching.py @@ -184,6 +184,8 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc] vllm_backend: Any | None = None, sym_tensor_indices: list[int] | None = None, aot_autograd_config: dict[str, Any] | None = None, + execution_code: str | None = None, + submod_names: list[str] | None = None, ) -> None: assert isinstance(graph_module, torch.fx.GraphModule) self.graph_module = graph_module @@ -194,6 +196,8 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc] self.shape_env = None self.vllm_backend = vllm_backend self.sym_tensor_indices = sym_tensor_indices + self.execution_code = execution_code + self.submod_names = submod_names self._fake_mode: Any | None = None import torch._functorch.config as functorch_config @@ -453,7 +457,7 @@ def reconstruct_serializable_fn_from_mega_artifact( standalone_compile_artifacts.load_all() - submod_names = standalone_compile_artifacts.submodule_names() + piecewise_submod_names = standalone_compile_artifacts.submodule_names() compiled_callables: dict[str, dict[str, Callable[..., Any]]] = {} for cache_key in standalone_compile_artifacts.submodule_bytes: @@ -473,13 +477,13 @@ def reconstruct_serializable_fn_from_mega_artifact( # spot check that cached submodules exist in the graph structure graph_children = {name for name, _ in split_gm.named_children()} - missing = set(submod_names) - graph_children + missing = set(piecewise_submod_names) - graph_children assert not missing, ( f"artifacts reference submodules not in graph: {missing}. " f"graph has: {sorted(graph_children)}" ) - for i, submod_name in enumerate(submod_names): + for i, submod_name in enumerate(piecewise_submod_names): assert submod_name in sym_shape_indices_map and submod_name in returns_tuple_map sym_shape_indices = sym_shape_indices_map[submod_name] @@ -490,7 +494,7 @@ def reconstruct_serializable_fn_from_mega_artifact( graph=None, # not needed for cached artifacts vllm_config=vllm_config, piecewise_compile_index=i, - total_piecewise_compiles=len(submod_names), + total_piecewise_compiles=len(piecewise_submod_names), sym_shape_indices=sym_shape_indices, vllm_backend=vllm_backend, returns_tuple=returns_tuple, @@ -498,7 +502,7 @@ def reconstruct_serializable_fn_from_mega_artifact( ) is_first = i == 0 - is_last = i == len(submod_names) - 1 + is_last = i == len(piecewise_submod_names) - 1 wrapped_backend = wrap_with_cudagraph_if_needed( piecewise_backend, vllm_config, @@ -513,6 +517,21 @@ def reconstruct_serializable_fn_from_mega_artifact( submod_name, ) + # Use codegen'd execution code if available, fall back to split_gm + execution_code = state.get("execution_code") + submod_names = state.get("submod_names") + if execution_code is not None and submod_names is not None: + from vllm.compilation.codegen import compile_execution_fn + + submod_callables = { + name: getattr(split_gm, name) for name, _ in split_gm.named_children() + } + runtime_callable = compile_execution_fn( + execution_code, submod_callables, submod_names + ) + else: + runtime_callable = split_gm + if compilation_config.cudagraph_copy_inputs: sym_tensor_indices = state["sym_tensor_indices"] input_buffers = [ @@ -521,9 +540,11 @@ def reconstruct_serializable_fn_from_mega_artifact( ) for idx in sym_tensor_indices ] - optimized_call = make_copy_and_call(sym_tensor_indices, input_buffers, split_gm) + optimized_call = make_copy_and_call( + sym_tensor_indices, input_buffers, runtime_callable + ) else: - optimized_call = split_gm + optimized_call = runtime_callable fn = VllmSerializableFunction( **state, diff --git a/vllm/compilation/codegen.py b/vllm/compilation/codegen.py new file mode 100644 index 00000000000..661b56cfb75 --- /dev/null +++ b/vllm/compilation/codegen.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Code generation for split_gm stitching graph execution. + +Generates a plain Python function that replaces the FX GraphModule's +interpreter-based execution of the stitching graph, eliminating +nn.Module.__call__ overhead and __getattr__ dispatch. +""" + +import operator +from collections.abc import Callable +from functools import partial +from typing import Any + +import torch.fx +from torch._dynamo.utils import dynamo_timed +from torch._logging import trace_structured + + +@dynamo_timed("vllm.generate_execution_code") +def generate_execution_code( + split_gm: torch.fx.GraphModule, +) -> tuple[str, list[str]]: + """Generate Python source code from a split_gm's stitching graph. + + Walks split_gm.graph.nodes and produces a function that calls + submodules via a __vllm_submods__ list, avoiding FX GraphModule overhead + and dict lookup cost. + + Args: + split_gm: The split graph module produced by split_graph(). + + Returns: + A tuple of (code, submod_names) where code is the Python source + and submod_names is the ordered list of submodule target names + corresponding to list indices used in the generated code. + """ + lines: list[str] = [] + param_names: list[str] = [] + submod_names: list[str] = [] + submod_index: dict[str, int] = {} + + # Build node ordering for liveness analysis. + nodes = list(split_gm.graph.nodes) + node_order = {node: i for i, node in enumerate(nodes)} + + # For each value-producing node, find the position of its last consumer. + # If the last consumer is the output node, skip (return handles cleanup). + # Otherwise, schedule a del after that consumer to free memory early. + del_after: dict[int, list[str]] = {} # position -> names to delete + for node in nodes: + if node.op == "output": + continue + users = list(node.users.keys()) + if not users: + continue + last_user = max(users, key=lambda u: node_order[u]) + if last_user.op == "output": + continue + del_after.setdefault(node_order[last_user], []).append(node.name) + + for i, node in enumerate(nodes): + if node.op == "placeholder": + param_names.append(node.name) + + elif node.op == "call_module": + target = node.target + if target not in submod_index: + submod_index[target] = len(submod_names) + submod_names.append(target) + idx = submod_index[target] + args_str = ", ".join(_node_ref(a) for a in node.args) + kwargs_str = ", ".join( + f"{k}={_node_ref(v)}" for k, v in node.kwargs.items() + ) + all_args = ", ".join(filter(None, [args_str, kwargs_str])) + lines.append(f" {node.name} = __vllm_submods__[{idx}]({all_args})") + + elif node.op == "call_function" and node.target is operator.getitem: + source = _node_ref(node.args[0]) + index = node.args[1] + assert isinstance(index, int) + lines.append(f" {node.name} = {source}[{index}]") + + elif node.op == "output": + assert len(node.args) == 1 + ret = _node_ref(node.args[0]) + lines.append(f" return {ret}") + + else: + raise RuntimeError(f"Unsupported node from codegen: {node.format_node()}") + + # Emit del for variables whose last use was this node. + if i in del_after: + names = sorted(del_after[i]) + lines.append(f" del {', '.join(names)}") + + assert len(param_names) > 0 + params = ", ".join(param_names) + header = f"def execution_fn({params}, *, __vllm_submods__):" + return "import torch\n" + "\n".join([header] + lines) + "\n", submod_names + + +@dynamo_timed("vllm.compile_execution_fn") +def compile_execution_fn( + code: str, + submod_callables: dict[str, Callable[..., Any]], + submod_names: list[str], +) -> Callable[..., Any]: + """Compile execution code and bind submodule callables. + + Args: + code: Python source from generate_execution_code(). + submod_callables: Mapping of submodule names to their callables. + submod_names: Ordered list of submodule names matching the indices + used in the generated code. + + Returns: + A callable that executes the stitching logic. + """ + trace_structured( + "artifact", + metadata_fn=lambda: { + "name": "vllm_execution_code", + "encoding": "string", + }, + payload_fn=lambda: code, + ) + namespace: dict[str, Any] = {} + exec(code, namespace) # noqa: S102 + fn = namespace["execution_fn"] + # Use .forward() directly to avoid nn.Module.__call__ overhead. + submods_list = [ + c.forward if isinstance(c, torch.fx.GraphModule) else c + for c in (submod_callables[name] for name in submod_names) + ] + return partial(fn, __vllm_submods__=submods_list) + + +def _node_ref(arg: Any) -> str: + """Convert an FX node argument to a source code reference recursively.""" + if isinstance(arg, torch.fx.Node): + return arg.name + if isinstance(arg, list): + return f"[{', '.join(_node_ref(x) for x in arg)}]" + if isinstance(arg, tuple): + items = ", ".join(_node_ref(x) for x in arg) + return f"({items},)" if len(arg) == 1 else f"({items})" + if isinstance(arg, dict): + return ( + "{" + + ", ".join(f"{_node_ref(k)}: {_node_ref(v)}" for k, v in arg.items()) + + "}" + ) + return repr(arg) From c0722f22de7159c7ed91fa8268bcecd87b35fe9e Mon Sep 17 00:00:00 2001 From: Julien Denize <40604584+juliendenize@users.noreply.github.com> Date: Thu, 16 Apr 2026 06:05:04 +0200 Subject: [PATCH 040/150] [Mistral Grammar] Fix tool and reasoning parsing (#39217) Signed-off-by: juliendenize --- .../tool_parsers/test_mistral_tool_parser.py | 938 ++++++++++++++---- .../mistral/test_mistral_tool_calls.py | 483 ++++++++- tests/tool_use/mistral/utils.py | 34 +- .../openai/chat_completion/protocol.py | 12 +- .../openai/chat_completion/serving.py | 69 +- vllm/entrypoints/openai/engine/serving.py | 34 +- vllm/entrypoints/serve/render/serving.py | 14 +- vllm/sampling_params.py | 14 +- vllm/tokenizers/mistral.py | 87 +- vllm/tool_parsers/mistral_tool_parser.py | 188 +++- 10 files changed, 1604 insertions(+), 269 deletions(-) diff --git a/tests/tool_parsers/test_mistral_tool_parser.py b/tests/tool_parsers/test_mistral_tool_parser.py index 064ccb39ef4..473eb716266 100644 --- a/tests/tool_parsers/test_mistral_tool_parser.py +++ b/tests/tool_parsers/test_mistral_tool_parser.py @@ -3,6 +3,7 @@ import json from collections.abc import Generator +from typing import Any from unittest.mock import MagicMock, patch import partial_json_parser @@ -23,24 +24,33 @@ from mistral_common.protocol.instruct.tool_calls import ( ToolChoiceEnum as MistralToolChoiceEnum, ) from partial_json_parser.core.options import Allow +from pydantic import ValidationError from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, DeltaMessage, DeltaToolCall, + ExtractedToolCallInformation, StructuralTagResponseFormat, ) +from vllm.entrypoints.openai.engine.protocol import FunctionCall as VllmFunctionCall +from vllm.reasoning.mistral_reasoning_parser import MistralReasoningParser from vllm.sampling_params import StructuredOutputsParams from vllm.tokenizers import TokenizerLike, get_tokenizer from vllm.tokenizers.detokenizer_utils import detokenize_incrementally from vllm.tokenizers.mistral import MistralTokenizer from vllm.tool_parsers.mistral_tool_parser import ( _DEFAULT_JSON_SCHEMA, + MistralStreamingResult, + MistralToolCall, MistralToolParser, ) +_DUMMY_REQUEST = ChatCompletionRequest(messages=[], model="test") + @pytest.fixture(scope="module") def mistral_pre_v11_tokenizer(): @@ -205,7 +215,7 @@ def stream_delta_message_generator( previous_token_ids, current_token_ids, delta_token_ids, - request=None, # type: ignore[arg-type] + request=_DUMMY_REQUEST, ) if delta_message: yield delta_message @@ -218,14 +228,18 @@ def stream_delta_message_generator( read_offset = new_read_offset -def test_extract_tool_calls_no_tools(mistral_pre_v11_tool_parser): +@pytest.mark.parametrize( + "parser_fixture", + ["mistral_pre_v11_tool_parser", "mistral_tool_parser"], + ids=["pre_v11", "v11"], +) +def test_extract_tool_calls_no_tools(parser_fixture, request): + parser = request.getfixturevalue(parser_fixture) model_output = "This is a test" - extracted_tool_calls = mistral_pre_v11_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - assert not extracted_tool_calls.tools_called - assert extracted_tool_calls.tool_calls == [] - assert extracted_tool_calls.content == model_output + result = parser.extract_tool_calls(model_output, request=_DUMMY_REQUEST) + assert result == ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) @pytest.mark.parametrize( @@ -234,6 +248,8 @@ def test_extract_tool_calls_no_tools(mistral_pre_v11_tool_parser): "single_tool_weather", "argument_before_name", "argument_before_name_and_name_in_argument", + "multiple_tools", + "content_before_tool", ], argnames=["model_output", "expected_tool_calls", "expected_content"], argvalues=[ @@ -292,14 +308,44 @@ def test_extract_tool_calls_no_tools(mistral_pre_v11_tool_parser): ], None, ), + ( + """[TOOL_CALLS] [{"name": "add", "arguments": {"a": 3.5, "b": 4}}, {"name": "get_current_weather", "arguments":{"city": "San Francisco", "state": "CA", "unit": "celsius"}}]""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="add", arguments=json.dumps({"a": 3.5, "b": 4}) + ) + ), + ToolCall( + function=FunctionCall( + name="get_current_weather", + arguments=json.dumps( + {"city": "San Francisco", "state": "CA", "unit": "celsius"} + ), + ) + ), + ], + None, + ), + ( + """Hello[TOOL_CALLS] [{"name": "add", "arguments":{"a": 1, "b": 2}}]""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="add", arguments=json.dumps({"a": 1, "b": 2}) + ) + ) + ], + "Hello", + ), ], ) def test_extract_tool_calls_pre_v11_tokenizer( mistral_pre_v11_tool_parser, model_output, expected_tool_calls, expected_content ): extracted_tool_calls = mistral_pre_v11_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] + model_output, request=_DUMMY_REQUEST + ) assert extracted_tool_calls.tools_called assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls) @@ -307,6 +353,46 @@ def test_extract_tool_calls_pre_v11_tokenizer( assert extracted_tool_calls.content == expected_content +def test_extract_tool_calls_pre_v11_multiple_bot_tokens_raises( + mistral_pre_v11_tool_parser, +): + model_output = ( + '[TOOL_CALLS] [{"name": "add", "arguments":{"a": 1}}]' + '[TOOL_CALLS] [{"name": "sub", "arguments":{"b": 2}}]' + ) + with pytest.raises(ValueError, match="Only one BOT token"): + mistral_pre_v11_tool_parser.extract_tool_calls( + model_output, request=_DUMMY_REQUEST + ) + + +def test_extract_tool_calls_pre_v11_regex_fallback_raises( + mistral_pre_v11_tool_parser, +): + """The regex fallback path finds valid JSON but does not re-serialize + the `arguments` dict to a string, causing a Pydantic + `ValidationError` when constructing `FunctionCall`.""" + model_output = ( + '[TOOL_CALLS] junk [{"name": "add", "arguments":{"a": 1, "b": 2}}] trail' + ) + with pytest.raises(ValidationError): + mistral_pre_v11_tool_parser.extract_tool_calls( + model_output, request=_DUMMY_REQUEST + ) + + +def test_extract_tool_calls_pre_v11_regex_fallback_fails( + mistral_pre_v11_tool_parser, +): + model_output = "[TOOL_CALLS] not json at all" + result = mistral_pre_v11_tool_parser.extract_tool_calls( + model_output, request=_DUMMY_REQUEST + ) + assert result == ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content="not json at all" + ) + + @pytest.mark.parametrize( ids=[ "single_tool_add", @@ -395,8 +481,8 @@ def test_extract_tool_calls( mistral_tool_parser, model_output, expected_tool_calls, expected_content ): extracted_tool_calls = mistral_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] + model_output, request=_DUMMY_REQUEST + ) assert extracted_tool_calls.tools_called assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls) @@ -404,6 +490,16 @@ def test_extract_tool_calls( assert extracted_tool_calls.content == expected_content +def test_extract_tool_calls_v11_without_args_skipped(mistral_tool_parser): + model_output = "[TOOL_CALLS]toolname_no_args" + result = mistral_tool_parser.extract_tool_calls( + model_output, request=_DUMMY_REQUEST + ) + assert result == ExtractedToolCallInformation( + tools_called=True, tool_calls=[], content=None + ) + + def _test_extract_tool_calls_streaming( tool_parser, tokenizer, model_output, tools, expected_tool_calls, expected_content ): @@ -669,17 +765,65 @@ def test_extract_tool_calls_streaming( ) +def test_extract_tool_calls_streaming_v11_no_tools( + mistral_tool_parser, mistral_tokenizer +): + model_output = "This is a test" + if isinstance(mistral_tokenizer, MistralTokenizer): + all_token_ids = mistral_tokenizer.encode(model_output) + else: + all_token_ids = mistral_tokenizer.encode(model_output, add_special_tokens=False) + skip_special = isinstance(mistral_tokenizer, MistralTokenizer) + collected_content = "" + previous_text = "" + previous_tokens = None + prefix_offset = 0 + read_offset = 0 + for i in range(len(all_token_ids)): + current_token_ids = all_token_ids[: i + 1] + previous_token_ids = all_token_ids[:i] + delta_token_ids = [all_token_ids[i]] + + new_tokens, delta_text, prefix_offset, read_offset = detokenize_incrementally( + tokenizer=mistral_tokenizer, + all_input_ids=current_token_ids, + prev_tokens=previous_tokens, + prefix_offset=prefix_offset, + read_offset=read_offset, + skip_special_tokens=skip_special, + spaces_between_special_tokens=True, + ) + current_text = previous_text + delta_text + previous_tokens = ( + previous_tokens + new_tokens if previous_tokens else new_tokens + ) + + delta_message = mistral_tool_parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=previous_token_ids, + current_token_ids=current_token_ids, + delta_token_ids=delta_token_ids, + request=_DUMMY_REQUEST, + ) + if delta_message and delta_message.content: + collected_content += delta_message.content + if delta_message: + assert not delta_message.tool_calls + + previous_text = current_text + + assert collected_content == model_output + + @pytest.mark.parametrize( - ids=[ - "single_tool_add", - "single_tool_weather", - "multiple_tool_calls", - "content_before_tool", - "complex", - ], - argnames=["model_output", "expected_tool_calls", "expected_content"], - argvalues=[ - ( + "parser_fixture, tokenizer_fixture, model_output," + " expected_tool_calls, expected_content", + [ + pytest.param( + "mistral_tool_parser", + "mistral_tokenizer", """[TOOL_CALLS]add_this_and_that{"a": 3.5, "b": 4}""", # noqa: E501 [ ToolCall( @@ -690,8 +834,11 @@ def test_extract_tool_calls_streaming( ) ], "", + id="v11-single_tool_add", ), - ( + pytest.param( + "mistral_tool_parser", + "mistral_tokenizer", """[TOOL_CALLS]get_current_weather{"city": "San Francisco", "state": "CA", "unit": "celsius"}""", # noqa: E501 [ ToolCall( @@ -704,8 +851,11 @@ def test_extract_tool_calls_streaming( ) ], "", + id="v11-single_tool_weather", ), - ( + pytest.param( + "mistral_tool_parser", + "mistral_tokenizer", """[TOOL_CALLS]add{"a": 3.5, "b": 4}[TOOL_CALLS]multiply{"a": 3, "b": 6}""", # noqa: E501 [ ToolCall( @@ -720,9 +870,11 @@ def test_extract_tool_calls_streaming( ), ], "", + id="v11-multiple_tool_calls", ), - ( - # Additional content should not be after the tool calls + pytest.param( + "mistral_tool_parser", + "mistral_tokenizer", """bla[TOOL_CALLS]add_this_and_that{"a": 3.5, "b": 4}""", # noqa: E501 [ ToolCall( @@ -733,9 +885,11 @@ def test_extract_tool_calls_streaming( ) ], "bla", + id="v11-content_before_tool", ), - ( - # Complex + pytest.param( + "mistral_tool_parser", + "mistral_tokenizer", """hi{hi[TOOL_CALLS]bash{"command": "print(\\"hello world!\\")\\nre.compile(r\'{}\')"}""", # noqa: E501 [ ToolCall( @@ -748,58 +902,19 @@ def test_extract_tool_calls_streaming( ) ], "hi{hi", + id="v11-complex", ), - ], -) -def test_extract_tool_calls_streaming_one_chunk( - mistral_tool_parser, - mistral_tokenizer, - model_output, - expected_tool_calls, - expected_content, -): - if isinstance(mistral_tokenizer, MistralTokenizer): - all_token_ids = mistral_tokenizer.encode(model_output) - else: - all_token_ids = mistral_tokenizer.encode(model_output, add_special_tokens=False) - all_token_ids = fix_tool_call_tokenization( - all_token_ids, mistral_tool_parser, mistral_tokenizer - ) - - delta_message = mistral_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=model_output, - delta_text=model_output, - previous_token_ids=[], - current_token_ids=all_token_ids, - delta_token_ids=all_token_ids, - request=None, - ) # type: ignore[arg-type] - assert isinstance(delta_message, DeltaMessage) - assert len(delta_message.tool_calls) == len(expected_tool_calls) - - assert_tool_calls(delta_message.tool_calls, expected_tool_calls) - - if delta_message.content is None: - assert expected_content == "" - else: - assert delta_message.content == expected_content - - -@pytest.mark.parametrize( - ids=[ - "no_tools", - "single_tool_add", - "single_tool_add_strings", - "single_tool_weather", - "argument_before_name", - "argument_before_name_and_name_in_argument", - "multiple_tools", - ], - argnames=["model_output", "expected_tool_calls", "expected_content"], - argvalues=[ - ("""This is a test""", [], """This is a test"""), - ( + pytest.param( + "mistral_pre_v11_tool_parser", + "mistral_pre_v11_tokenizer", + """This is a test""", + [], + """This is a test""", + id="pre_v11-no_tools", + ), + pytest.param( + "mistral_pre_v11_tool_parser", + "mistral_pre_v11_tokenizer", """[TOOL_CALLS] [ {"name":"add" , "arguments" : {"a": 3, "b": 4} } ]""", # noqa: E501 [ ToolCall( @@ -809,8 +924,11 @@ def test_extract_tool_calls_streaming_one_chunk( ) ], "", + id="pre_v11-single_tool_add", ), - ( + pytest.param( + "mistral_pre_v11_tool_parser", + "mistral_pre_v11_tokenizer", """[TOOL_CALLS] [{"name": "add", "arguments":{"a": "3", "b": "4"}}]""", # noqa: E501 [ ToolCall( @@ -820,8 +938,11 @@ def test_extract_tool_calls_streaming_one_chunk( ) ], "", + id="pre_v11-single_tool_add_strings", ), - ( + pytest.param( + "mistral_pre_v11_tool_parser", + "mistral_pre_v11_tokenizer", """[TOOL_CALLS] [{"name": "get_current_weather", "arguments": {"city": "San Francisco", "state": "CA", "unit": "celsius"}}]""", # noqa: E501 [ ToolCall( @@ -834,8 +955,11 @@ def test_extract_tool_calls_streaming_one_chunk( ) ], "", + id="pre_v11-single_tool_weather", ), - ( + pytest.param( + "mistral_pre_v11_tool_parser", + "mistral_pre_v11_tokenizer", """[TOOL_CALLS] [{"arguments": {"city": "San Francisco", "state": "CA", "unit": "celsius"}, "name": "get_current_weather"}]""", # noqa: E501 [ ToolCall( @@ -848,8 +972,11 @@ def test_extract_tool_calls_streaming_one_chunk( ) ], "", + id="pre_v11-argument_before_name", ), - ( + pytest.param( + "mistral_pre_v11_tool_parser", + "mistral_pre_v11_tokenizer", """[TOOL_CALLS] [{"arguments": {"name": "John Doe"}, "name": "get_age"}]""", # noqa: E501 [ ToolCall( @@ -864,8 +991,11 @@ def test_extract_tool_calls_streaming_one_chunk( ) ], "", + id="pre_v11-argument_before_name_and_name_in_argument", ), - ( + pytest.param( + "mistral_pre_v11_tool_parser", + "mistral_pre_v11_tokenizer", """[TOOL_CALLS] [{"arguments": {"a": 3.5, "b": 4}, "name": "add"}, {"arguments":{"city": "San Francisco", "state": "CA", "unit": "celsius"}, "name": "get_current_weather"}]""", # noqa: E501 [ ToolCall( @@ -883,35 +1013,50 @@ def test_extract_tool_calls_streaming_one_chunk( ), ], "", + id="pre_v11-multiple_tools", + ), + pytest.param( + "mistral_pre_v11_tool_parser", + "mistral_pre_v11_tokenizer", + """Some text[TOOL_CALLS] [{"name": "add", "arguments":{"a": 1, "b": 2}}]""", # noqa: E501 + [ + ToolCall( + function=FunctionCall( + name="add", arguments=json.dumps({"a": 1, "b": 2}) + ) + ) + ], + "Some text", + id="pre_v11-content_before_tool", ), ], ) -def test_extract_tool_calls_streaming_pre_v11_tokenizer_one_chunk( - mistral_pre_v11_tool_parser, - mistral_pre_v11_tokenizer, +def test_extract_tool_calls_streaming_one_chunk( + parser_fixture, + tokenizer_fixture, model_output, expected_tool_calls, expected_content, + request, ): - if isinstance(mistral_pre_v11_tokenizer, MistralTokenizer): - all_token_ids = mistral_pre_v11_tokenizer.encode(model_output) - else: - all_token_ids = mistral_pre_v11_tokenizer.encode( - model_output, add_special_tokens=False - ) - all_token_ids = fix_tool_call_tokenization( - all_token_ids, mistral_pre_v11_tool_parser, mistral_pre_v11_tokenizer - ) + tool_parser = request.getfixturevalue(parser_fixture) + tokenizer = request.getfixturevalue(tokenizer_fixture) - delta_message = mistral_pre_v11_tool_parser.extract_tool_calls_streaming( + if isinstance(tokenizer, MistralTokenizer): + all_token_ids = tokenizer.encode(model_output) + else: + all_token_ids = tokenizer.encode(model_output, add_special_tokens=False) + all_token_ids = fix_tool_call_tokenization(all_token_ids, tool_parser, tokenizer) + + delta_message = tool_parser.extract_tool_calls_streaming( previous_text="", current_text=model_output, delta_text=model_output, previous_token_ids=[], current_token_ids=all_token_ids, delta_token_ids=all_token_ids, - request=None, - ) # type: ignore[arg-type] + request=_DUMMY_REQUEST, + ) assert isinstance(delta_message, DeltaMessage) assert len(delta_message.tool_calls) == len(expected_tool_calls) @@ -923,67 +1068,107 @@ def test_extract_tool_calls_streaming_pre_v11_tokenizer_one_chunk( assert delta_message.content == expected_content -def test_fast_detokenization_text_detection(mistral_tool_parser): - """Regression: bot_token in text but not token_ids (PR #37209).""" - model_output = '[TOOL_CALLS]add{"a": 1, "b": 2}' - # Token IDs that do NOT contain bot_token_id. - fake_token_ids = list(range(99, 99 + 20)) - - # First delta: pure content, no bot token yet - delta_message_before = mistral_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text="Hello", - delta_text="Hello", - previous_token_ids=[], - current_token_ids=[99], - delta_token_ids=[99], - request=None, - ) - assert delta_message_before is not None - assert delta_message_before.content == "Hello" - assert not delta_message_before.tool_calls - - # Second delta: bot token in text but NOT in token_ids - delta_message = mistral_tool_parser.extract_tool_calls_streaming( - previous_text="Hello", - current_text="Hello" + model_output, - delta_text=model_output, - previous_token_ids=[99], - current_token_ids=fake_token_ids, - delta_token_ids=fake_token_ids[1:], - request=None, - ) - assert delta_message is not None - assert delta_message.tool_calls is not None - assert len(delta_message.tool_calls) > 0 - assert delta_message.tool_calls[0].function is not None - assert delta_message.tool_calls[0].function.name == "add" - - -def test_fast_detokenization_text_detection_pre_v11( - mistral_pre_v11_tool_parser, +@pytest.mark.parametrize( + "parser_fixture, model_output, fake_count, two_phase", + [ + pytest.param( + "mistral_tool_parser", + '[TOOL_CALLS]add{"a": 1, "b": 2}', + 20, + True, + id="v11", + ), + pytest.param( + "mistral_pre_v11_tool_parser", + '[TOOL_CALLS] [{"name": "add", "arguments":{"a": 1, "b": 2}}]', + 30, + False, + id="pre_v11", + ), + ], +) +def test_fast_detokenization_text_detection( + parser_fixture, model_output, fake_count, two_phase, request ): - """Regression: bot_token text detection for pre-v11 tokenizer (PR #37209).""" - model_output = '[TOOL_CALLS] [{"name": "add", "arguments":{"a": 1, "b": 2}}]' + """Regression: bot_token in text but not token_ids (PR #37209).""" + parser = request.getfixturevalue(parser_fixture) + # Token IDs that do NOT contain bot_token_id. + fake_token_ids = list(range(99, 99 + fake_count)) - fake_token_ids = list(range(99, 99 + 30)) + if two_phase: + # First delta: pure content, no bot token yet + delta_message_before = parser.extract_tool_calls_streaming( + previous_text="", + current_text="Hello", + delta_text="Hello", + previous_token_ids=[], + current_token_ids=[99], + delta_token_ids=[99], + request=_DUMMY_REQUEST, + ) + assert delta_message_before is not None + assert delta_message_before.content == "Hello" + assert not delta_message_before.tool_calls - delta_message = mistral_pre_v11_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=model_output, + previous_text = "Hello" + current_text = "Hello" + model_output + previous_token_ids = [99] + delta_token_ids = fake_token_ids[1:] + else: + previous_text = "" + current_text = model_output + previous_token_ids = [] + delta_token_ids = fake_token_ids + + delta_message = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, delta_text=model_output, - previous_token_ids=[], + previous_token_ids=previous_token_ids, current_token_ids=fake_token_ids, - delta_token_ids=fake_token_ids, - request=None, + delta_token_ids=delta_token_ids, + request=_DUMMY_REQUEST, ) assert delta_message is not None assert delta_message.tool_calls is not None - assert len(delta_message.tool_calls) > 0 + assert len(delta_message.tool_calls) == 1 assert delta_message.tool_calls[0].function is not None assert delta_message.tool_calls[0].function.name == "add" +@pytest.mark.parametrize( + "parser_fixture, patched_method, current_text", + [ + ( + "mistral_tool_parser", + "_extract_tool_calls_streaming", + "[TOOL_CALLS]add{}", + ), + ( + "mistral_pre_v11_tool_parser", + "_extract_tool_calls_streaming_pre_v11_tokenizer", + '[TOOL_CALLS] [{"name":"a","arguments":{}}]', + ), + ], + ids=["v11", "pre_v11"], +) +def test_extract_tool_calls_streaming_exception_returns_none( + parser_fixture, patched_method, current_text, request +): + parser = request.getfixturevalue(parser_fixture) + with patch.object(parser, patched_method, side_effect=RuntimeError("boom")): + result = parser.extract_tool_calls_streaming( + previous_text="", + current_text=current_text, + delta_text=current_text, + previous_token_ids=[], + current_token_ids=[parser.bot_token_id], + delta_token_ids=[parser.bot_token_id], + request=_DUMMY_REQUEST, + ) + assert result is None + + SAMPLE_TOOLS_DICTS = [ { "type": "function", @@ -1238,57 +1423,444 @@ def test_adjust_request_response_format_generates_grammar( assert len(result.structured_outputs.grammar) > 0 -def test_adjust_request_tool_choice_none_with_json_schema_uses_json_schema_factory( +@pytest.mark.parametrize( + "tool_choice, expected_method, not_called_method", + [ + ("none", "get_lark_for_json_schema", None), + ("auto", "get_lark_from_jinja", "get_lark_for_json_schema"), + ], + ids=["none_uses_json_schema_factory", "auto_uses_jinja_factory"], +) +def test_adjust_request_tool_choice_with_json_schema_factory_routing( mistral_tool_parser: MistralToolParser, + tool_choice: str, + expected_method: str, + not_called_method: str | None, ) -> None: request = _make_request( - tool_choice="none", + tool_choice=tool_choice, structured_outputs=StructuredOutputsParams(json='{"type": "object"}'), ) factory = mistral_tool_parser.model_tokenizer.grammar_factory - with patch.object( - factory, - "get_lark_for_json_schema", - wraps=factory.get_lark_for_json_schema, - ) as mock_json_schema: - result = mistral_tool_parser.adjust_request(request) + patches = { + expected_method: patch.object( + factory, + expected_method, + wraps=getattr(factory, expected_method), + ), + } + if not_called_method: + patches[not_called_method] = patch.object( + factory, + not_called_method, + wraps=getattr(factory, not_called_method), + ) - mock_json_schema.assert_called_once() - assert mock_json_schema.call_args.kwargs["json_schema"] == {"type": "object"} + with patches[expected_method] as mock_expected: + ctx = patches[not_called_method] if not_called_method else None + if ctx: + with ctx as mock_not_called: + result = mistral_tool_parser.adjust_request(request) + mock_not_called.assert_not_called() + else: + result = mistral_tool_parser.adjust_request(request) + + mock_expected.assert_called_once() + assert mock_expected.call_args.kwargs["json_schema"] == {"type": "object"} assert result.structured_outputs is not None assert isinstance(result.structured_outputs.grammar, str) assert len(result.structured_outputs.grammar) > 0 -def test_adjust_request_tool_choice_auto_with_json_schema_uses_jinja_factory( +def test_grammar_from_tool_parser_default_false() -> None: + request = _make_request() + assert request._grammar_from_tool_parser is False + + +def test_grammar_from_tool_parser_set_by_adjust_request( mistral_tool_parser: MistralToolParser, ) -> None: - request = _make_request( - tool_choice="auto", - structured_outputs=StructuredOutputsParams(json='{"type": "object"}'), - ) - factory = mistral_tool_parser.model_tokenizer.grammar_factory + request = _make_request() + result = mistral_tool_parser.adjust_request(request) + assert result._grammar_from_tool_parser is True - with ( - patch.object( - factory, - "get_lark_for_json_schema", - wraps=factory.get_lark_for_json_schema, - ) as mock_json_schema, - patch.object( - factory, - "get_lark_from_jinja", - wraps=factory.get_lark_from_jinja, - ) as mock_jinja, - ): - result = mistral_tool_parser.adjust_request(request) - mock_jinja.assert_called_once() - assert mock_jinja.call_args.kwargs["json_schema"] == {"type": "object"} - mock_json_schema.assert_not_called() +@pytest.mark.parametrize( + "tool_calls, expected_len", + [ + (None, 0), + ([], 0), + ([VllmFunctionCall(id="abc123xyz", name="f", arguments="{}")], 1), + ([VllmFunctionCall(name="f", arguments="{}")], 1), + ( + [ + VllmFunctionCall(id="fixed1234", name="a", arguments='{"x": 1}'), + VllmFunctionCall(name="b", arguments='{"y": 2}'), + ], + 2, + ), + ], + ids=["none", "empty", "with_id", "without_id", "mixed"], +) +def test_build_non_streaming_tool_calls( + tool_calls: list[VllmFunctionCall] | None, + expected_len: int, +) -> None: + result = MistralToolParser.build_non_streaming_tool_calls(tool_calls) + assert len(result) == expected_len - assert result.structured_outputs is not None - assert isinstance(result.structured_outputs.grammar, str) - assert len(result.structured_outputs.grammar) > 0 + if tool_calls is None: + return + + for i, tc in enumerate(result): + assert isinstance(tc, MistralToolCall) + assert tc.type == "function" + + input_tc = tool_calls[i] + if input_tc.id: + assert tc.id == input_tc.id + else: + assert len(tc.id) == 9 + assert tc.id.isalnum() + + assert tc.function.name == input_tc.name + assert tc.function.arguments == input_tc.arguments + + +class TestExtractMaybeReasoningAndToolStreaming: + r"""Tests for `MistralToolParser.extract_maybe_reasoning_and_tool_streaming`.""" + + @pytest.fixture + def parser(self) -> MistralToolParser: + mock_tokenizer = MagicMock() + mock_tokenizer.get_vocab.return_value = {"[TOOL_CALLS]": 1} + return MistralToolParser(mock_tokenizer) + + @pytest.fixture + def request_obj(self) -> ChatCompletionRequest: + return _make_request() + + @staticmethod + def _call( + parser: MistralToolParser, + request: ChatCompletionRequest, + *, + reasoning_parser: Any = None, + previous_text: str = "", + current_text: str = "hello", + delta_text: str = "hello", + previous_token_ids: list[int] | None = None, + current_token_ids: list[int] | None = None, + output_token_ids: list[int] | None = None, + reasoning_ended: bool = False, + prompt_is_reasoning_end: bool | None = None, + ) -> MistralStreamingResult: + return parser.extract_maybe_reasoning_and_tool_streaming( + reasoning_parser=reasoning_parser, + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=previous_token_ids or [], + current_token_ids=current_token_ids or [1, 2, 3], + output_token_ids=output_token_ids or [1, 2, 3], + reasoning_ended=reasoning_ended, + prompt_is_reasoning_end=prompt_is_reasoning_end, + request=request, + ) + + def test_no_reasoning_tools_called( + self, parser: MistralToolParser, request_obj: ChatCompletionRequest + ) -> None: + tool_delta = DeltaMessage( + tool_calls=[ + DeltaToolCall( + index=0, + function=DeltaFunctionCall(name="f", arguments="{}"), + ) + ] + ) + with patch.object( + parser, "extract_tool_calls_streaming", return_value=tool_delta + ): + result = self._call(parser, request_obj, reasoning_parser=None) + + assert result == MistralStreamingResult( + delta_message=tool_delta, + reasoning_ended=False, + tools_called=True, + current_text="hello", + current_token_ids=[1, 2, 3], + ) + + def test_no_reasoning_no_tools( + self, parser: MistralToolParser, request_obj: ChatCompletionRequest + ) -> None: + content_delta = DeltaMessage(content="hello") + with patch.object( + parser, "extract_tool_calls_streaming", return_value=content_delta + ): + result = self._call(parser, request_obj, reasoning_parser=None) + + assert result == MistralStreamingResult( + delta_message=content_delta, + reasoning_ended=False, + tools_called=False, + current_text="hello", + current_token_ids=[1, 2, 3], + ) + + def test_mistral_reasoning_parser_no_think_token( + self, parser: MistralToolParser, request_obj: ChatCompletionRequest + ) -> None: + mock_rp = MagicMock(spec=MistralReasoningParser) + mock_rp.start_token_id = 999 + content_delta = DeltaMessage(content="direct") + with patch.object( + parser, "extract_tool_calls_streaming", return_value=content_delta + ): + result = self._call( + parser, + request_obj, + reasoning_parser=mock_rp, + reasoning_ended=False, + current_token_ids=[1, 2, 3], + ) + + mock_rp.extract_reasoning_streaming.assert_not_called() + assert result == MistralStreamingResult( + delta_message=content_delta, + reasoning_ended=False, + tools_called=False, + current_text="hello", + current_token_ids=[1, 2, 3], + ) + + def test_mistral_reasoning_parser_with_think_token( + self, parser: MistralToolParser, request_obj: ChatCompletionRequest + ) -> None: + mock_rp = MagicMock(spec=MistralReasoningParser) + mock_rp.start_token_id = 999 + mock_rp.extract_reasoning_streaming.return_value = DeltaMessage( + reasoning="thinking..." + ) + mock_rp.is_reasoning_end_streaming.return_value = False + + result = self._call( + parser, + request_obj, + reasoning_parser=mock_rp, + reasoning_ended=False, + current_token_ids=[1, 999, 3], + ) + + mock_rp.extract_reasoning_streaming.assert_called_once() + assert result == MistralStreamingResult( + delta_message=DeltaMessage(reasoning="thinking..."), + reasoning_ended=False, + tools_called=False, + current_text="hello", + current_token_ids=[1, 999, 3], + ) + + def test_non_mistral_reasoning_parser_always_expects_thinking( + self, parser: MistralToolParser, request_obj: ChatCompletionRequest + ) -> None: + mock_rp = MagicMock() + mock_rp.start_token_id = 999 + mock_rp.extract_reasoning_streaming.return_value = DeltaMessage( + reasoning="thinking..." + ) + mock_rp.is_reasoning_end_streaming.return_value = False + + result = self._call( + parser, + request_obj, + reasoning_parser=mock_rp, + reasoning_ended=False, + current_token_ids=[1, 2, 3], + ) + + mock_rp.extract_reasoning_streaming.assert_called_once() + assert result == MistralStreamingResult( + delta_message=DeltaMessage(reasoning="thinking..."), + reasoning_ended=False, + tools_called=False, + current_text="hello", + current_token_ids=[1, 2, 3], + ) + + def test_reasoning_already_ended_no_reset( + self, parser: MistralToolParser, request_obj: ChatCompletionRequest + ) -> None: + content_delta = DeltaMessage(content="content") + with patch.object( + parser, "extract_tool_calls_streaming", return_value=content_delta + ) as mock_extract: + result = self._call( + parser, + request_obj, + reasoning_parser=MagicMock(), + reasoning_ended=True, + previous_text="prior_tool_text", + previous_token_ids=[10, 20], + current_text="prior_tool_texthello", + current_token_ids=[10, 20, 1, 2, 3], + ) + + _, call_kwargs = mock_extract.call_args + assert call_kwargs["previous_text"] == "prior_tool_text" + assert call_kwargs["previous_token_ids"] == [10, 20] + + assert result == MistralStreamingResult( + delta_message=content_delta, + reasoning_ended=True, + tools_called=False, + current_text="prior_tool_texthello", + current_token_ids=[10, 20, 1, 2, 3], + ) + + def test_pre_v15_ignores_prompt_reasoning_end( + self, parser: MistralToolParser, request_obj: ChatCompletionRequest + ) -> None: + mock_tokenizer = MagicMock(spec=MistralTokenizer) + mock_tokenizer.version = 13 + parser.model_tokenizer = mock_tokenizer + + mock_rp = MagicMock(spec=MistralReasoningParser) + mock_rp.start_token_id = 999 + mock_rp.extract_reasoning_streaming.return_value = DeltaMessage( + reasoning="thinking..." + ) + mock_rp.is_reasoning_end_streaming.return_value = False + + result = self._call( + parser, + request_obj, + reasoning_parser=mock_rp, + reasoning_ended=False, + prompt_is_reasoning_end=True, + current_token_ids=[999, 1, 2], + ) + + mock_rp.extract_reasoning_streaming.assert_called_once() + assert result == MistralStreamingResult( + delta_message=DeltaMessage(reasoning="thinking..."), + reasoning_ended=False, + tools_called=False, + current_text="hello", + current_token_ids=[999, 1, 2], + ) + + def test_non_pre_v15_prompt_reasoning_end( + self, parser: MistralToolParser, request_obj: ChatCompletionRequest + ) -> None: + mock_tokenizer = MagicMock(spec=MistralTokenizer) + mock_tokenizer.version = 15 + parser.model_tokenizer = mock_tokenizer + + mock_rp = MagicMock(spec=MistralReasoningParser) + mock_rp.start_token_id = 999 + + content_delta = DeltaMessage(content="after reasoning") + with patch.object( + parser, "extract_tool_calls_streaming", return_value=content_delta + ): + result = self._call( + parser, + request_obj, + reasoning_parser=mock_rp, + reasoning_ended=False, + prompt_is_reasoning_end=True, + current_token_ids=[999, 1, 2], + output_token_ids=[10, 20, 30], + ) + + mock_rp.extract_reasoning_streaming.assert_not_called() + assert result == MistralStreamingResult( + delta_message=content_delta, + reasoning_ended=True, + tools_called=False, + current_text="hello", + current_token_ids=[10, 20, 30], + ) + + def test_reasoning_end_transition_with_content( + self, parser: MistralToolParser, request_obj: ChatCompletionRequest + ) -> None: + """When reasoning ends and the delta has content, that content is + cleared from delta_message and used as current_text for tool parsing.""" + mock_rp = MagicMock() + mock_rp.start_token_id = 999 + mock_rp.extract_reasoning_streaming.return_value = DeltaMessage( + reasoning="think", content="leftover" + ) + mock_rp.is_reasoning_end_streaming.return_value = True + mock_rp.extract_content_ids.return_value = [50, 51] + + content_delta = DeltaMessage(content="leftover") + with patch.object( + parser, "extract_tool_calls_streaming", return_value=content_delta + ) as mock_extract: + result = self._call( + parser, + request_obj, + reasoning_parser=mock_rp, + reasoning_ended=False, + current_token_ids=[999, 1, 2], + output_token_ids=[10, 20, 30], + ) + + mock_rp.extract_content_ids.assert_called_once_with([10, 20, 30]) + _, call_kwargs = mock_extract.call_args + assert call_kwargs["previous_text"] == "" + assert call_kwargs["previous_token_ids"] == [] + assert call_kwargs["delta_text"] == "leftover" + assert call_kwargs["current_token_ids"] == [50, 51] + + assert result == MistralStreamingResult( + delta_message=content_delta, + reasoning_ended=True, + tools_called=False, + current_text="leftover", + current_token_ids=[50, 51], + ) + + def test_reasoning_end_transition_without_content( + self, parser: MistralToolParser, request_obj: ChatCompletionRequest + ) -> None: + """When reasoning ends but the delta has no content, current_text + is set to empty string.""" + mock_rp = MagicMock() + mock_rp.start_token_id = 999 + mock_rp.extract_reasoning_streaming.return_value = DeltaMessage( + reasoning="think" + ) + mock_rp.is_reasoning_end_streaming.return_value = True + mock_rp.extract_content_ids.return_value = [50, 51] + + empty_delta = DeltaMessage(content="") + with patch.object( + parser, "extract_tool_calls_streaming", return_value=empty_delta + ) as mock_extract: + result = self._call( + parser, + request_obj, + reasoning_parser=mock_rp, + reasoning_ended=False, + current_token_ids=[999, 1, 2], + output_token_ids=[10, 20, 30], + ) + + _, call_kwargs = mock_extract.call_args + assert call_kwargs["delta_text"] == "" + assert call_kwargs["current_token_ids"] == [50, 51] + + assert result == MistralStreamingResult( + delta_message=empty_delta, + reasoning_ended=True, + tools_called=False, + current_text="", + current_token_ids=[50, 51], + ) diff --git a/tests/tool_use/mistral/test_mistral_tool_calls.py b/tests/tool_use/mistral/test_mistral_tool_calls.py index 3c4a543abe4..6dcfd43a949 100644 --- a/tests/tool_use/mistral/test_mistral_tool_calls.py +++ b/tests/tool_use/mistral/test_mistral_tool_calls.py @@ -1,25 +1,198 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +from dataclasses import dataclass, field + import openai import pytest -from tests.tool_use.utils import MESSAGES_ASKING_FOR_TOOLS, WEATHER_TOOL +from tests.tool_use.utils import ( + MESSAGES_ASKING_FOR_PARALLEL_TOOLS, + MESSAGES_ASKING_FOR_TOOLS, + MESSAGES_WITH_TOOL_RESPONSE, + MESSAGES_WITHOUT_TOOLS, + SEARCH_TOOL, + SEED, + WEATHER_TOOL, + ensure_system_prompt, +) + +from .utils import ServerConfig + + +def _requires_tool_parser(server_config: ServerConfig) -> None: + r"""Skip test if server was not started with --tool-call-parser.""" + if "--tool-call-parser" not in server_config.get("arguments", []): + pytest.skip( + f"Skipping: {server_config['model']} not configured with --tool-call-parser" + ) + + +def _is_pre_v11(server_config: ServerConfig) -> bool: + r"""Pre-v11 Mistral models lack grammar-based tool call enforcement.""" + return "7B" in server_config.get("model", "") + + +@dataclass +class StreamedToolCallResult: + r"""Accumulated result from streaming a single tool call.""" + + function_name: str | None = None + function_args_str: str = "" + tool_call_id: str | None = None + role_name: str | None = None + finish_reason_count: int = 0 + finish_reason: str | None = None + + +async def _collect_streamed_tool_call( + stream: openai.AsyncStream, + *, + expected_finish_reason: str = "tool_calls", +) -> StreamedToolCallResult: + result = StreamedToolCallResult() + + async for chunk in stream: + if chunk.choices[0].finish_reason: + result.finish_reason_count += 1 + result.finish_reason = chunk.choices[0].finish_reason + assert chunk.choices[0].finish_reason == expected_finish_reason + + if chunk.choices[0].delta.role: + assert not result.role_name or result.role_name == "assistant" + result.role_name = "assistant" + + streamed_tool_calls = chunk.choices[0].delta.tool_calls + if streamed_tool_calls and len(streamed_tool_calls) > 0: + assert len(streamed_tool_calls) == 1 + tool_call = streamed_tool_calls[0] + + if tool_call.id: + assert not result.tool_call_id + result.tool_call_id = tool_call.id + + if tool_call.function: + if tool_call.function.name: + assert result.function_name is None + result.function_name = tool_call.function.name + if tool_call.function.arguments: + result.function_args_str += tool_call.function.arguments + + return result + + +@dataclass +class StreamedContentResult: + r"""Accumulated result from streaming a content-only response.""" + + chunks: list[str] = field(default_factory=list) + finish_reason_count: int = 0 + finish_reason: str | None = None + role_sent: bool = False + + +async def _collect_streamed_content( + stream: openai.AsyncStream, + *, + expected_finish_reason: str | None = None, + no_tool_calls: bool = True, +) -> StreamedContentResult: + r"""Consume a streaming response and collect text content.""" + result = StreamedContentResult() + + async for chunk in stream: + delta = chunk.choices[0].delta + + if delta.role: + assert not result.role_sent + assert delta.role == "assistant" + result.role_sent = True + + if delta.content: + result.chunks.append(delta.content) + + if chunk.choices[0].finish_reason is not None: + result.finish_reason_count += 1 + result.finish_reason = chunk.choices[0].finish_reason + if expected_finish_reason is not None: + assert result.finish_reason == expected_finish_reason + + if no_tool_calls: + assert not delta.tool_calls or len(delta.tool_calls) == 0 + + return result + + +@dataclass +class StreamedParallelToolCallResult: + r"""Accumulated result from streaming parallel tool calls.""" + + function_names: list[str] = field(default_factory=list) + function_args_strs: list[str] = field(default_factory=list) + tool_call_ids: list[str] = field(default_factory=list) + role_name: str | None = None + finish_reason_count: int = 0 + + +async def _collect_streamed_parallel_tool_calls( + stream: openai.AsyncStream, +) -> StreamedParallelToolCallResult: + r"""Consume a streaming response and collect parallel tool calls.""" + result = StreamedParallelToolCallResult() + tool_call_idx: int = -1 + + async for chunk in stream: + if chunk.choices[0].finish_reason: + result.finish_reason_count += 1 + assert chunk.choices[0].finish_reason == "tool_calls" + + if chunk.choices[0].delta.role: + assert not result.role_name or result.role_name == "assistant" + result.role_name = "assistant" + + streamed_tool_calls = chunk.choices[0].delta.tool_calls + if streamed_tool_calls and len(streamed_tool_calls) > 0: + assert len(streamed_tool_calls) == 1 + tool_call = streamed_tool_calls[0] + + if tool_call.index != tool_call_idx: + tool_call_idx = tool_call.index + result.function_args_strs.append("") + result.tool_call_ids.append("") + + if tool_call.id: + result.tool_call_ids[tool_call.index] = tool_call.id + + if tool_call.function: + if tool_call.function.name: + result.function_names.append(tool_call.function.name) + if tool_call.function.arguments: + result.function_args_strs[tool_call.index] += ( + tool_call.function.arguments + ) + + return result # test: a tool_choice with mistral-tokenizer results in an ID of length 9 @pytest.mark.asyncio -async def test_tool_call_with_tool_choice(client: openai.AsyncOpenAI): +async def test_tool_call_with_tool_choice( + client: openai.AsyncOpenAI, server_config: ServerConfig +) -> None: + _requires_tool_parser(server_config) + models = await client.models.list() model_name: str = models.data[0].id chat_completion = await client.chat.completions.create( - messages=MESSAGES_ASKING_FOR_TOOLS, + messages=ensure_system_prompt(MESSAGES_ASKING_FOR_TOOLS, server_config), temperature=0, max_completion_tokens=100, model=model_name, tools=[WEATHER_TOOL], tool_choice=WEATHER_TOOL, logprobs=False, + seed=SEED, ) choice = chat_completion.choices[0] @@ -28,3 +201,307 @@ async def test_tool_call_with_tool_choice(client: openai.AsyncOpenAI): assert choice.message.role == "assistant" assert choice.message.tool_calls is None or len(choice.message.tool_calls) == 1 assert len(choice.message.tool_calls[0].id) == 9 # length of 9 for mistral + + +_NOT_SET = object() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tools, tool_choice, streaming_id_len_pre_v11", + [ + pytest.param( + [WEATHER_TOOL, SEARCH_TOOL], + _NOT_SET, + 9, + id="auto", + ), + pytest.param( + [WEATHER_TOOL], + "required", + 30, + id="required", + ), + ], +) +async def test_tool_call_auto_or_required( + client: openai.AsyncOpenAI, + server_config: ServerConfig, + tools: list, + tool_choice: object, + streaming_id_len_pre_v11: int, +) -> None: + _requires_tool_parser(server_config) + + models = await client.models.list() + model_name: str = models.data[0].id + + create_kwargs: dict = { + "messages": ensure_system_prompt(MESSAGES_ASKING_FOR_TOOLS, server_config), + "temperature": 0, + "max_completion_tokens": 100, + "model": model_name, + "tools": tools, + "logprobs": False, + "seed": SEED, + } + if tool_choice is not _NOT_SET: + create_kwargs["tool_choice"] = tool_choice + + # --- non-streaming --- + chat_completion = await client.chat.completions.create(**create_kwargs) + + choice = chat_completion.choices[0] + tool_calls = choice.message.tool_calls + + assert choice.finish_reason == "tool_calls" + assert tool_calls is not None and len(tool_calls) >= 1 + assert tool_calls[0].function.name == "get_current_weather" + parsed_arguments = json.loads(tool_calls[0].function.arguments) + assert "city" in parsed_arguments + assert len(tool_calls[0].id) == 9 + + # --- streaming --- + stream = await client.chat.completions.create(**create_kwargs, stream=True) + + result = await _collect_streamed_tool_call(stream) + + assert result.finish_reason_count == 1 + assert result.role_name == "assistant" + assert result.function_name == "get_current_weather" + streamed_args = json.loads(result.function_args_str) + assert isinstance(result.tool_call_id, str) + if _is_pre_v11(server_config): + assert len(result.tool_call_id) == streaming_id_len_pre_v11 + else: + assert len(result.tool_call_id) == 9 + assert parsed_arguments == streamed_args + + +@pytest.mark.asyncio +async def test_tool_call_none_with_tools( + client: openai.AsyncOpenAI, server_config: ServerConfig +) -> None: + _requires_tool_parser(server_config) + + models = await client.models.list() + model_name: str = models.data[0].id + + # --- non-streaming --- + chat_completion = await client.chat.completions.create( + messages=ensure_system_prompt(MESSAGES_ASKING_FOR_TOOLS, server_config), + temperature=0, + max_completion_tokens=100, + model=model_name, + tools=[WEATHER_TOOL], + tool_choice="none", + logprobs=False, + seed=SEED, + ) + + choice = chat_completion.choices[0] + + assert choice.finish_reason != "tool_calls" + assert choice.message.tool_calls is None or len(choice.message.tool_calls) == 0 + assert choice.message.content is not None + # Without grammar enforcement, pre-v11 models may still emit [TOOL_CALLS] + if not _is_pre_v11(server_config): + assert "[TOOL_CALLS]" not in choice.message.content + + non_streaming_content = choice.message.content + + # --- streaming --- + stream = await client.chat.completions.create( + messages=ensure_system_prompt(MESSAGES_ASKING_FOR_TOOLS, server_config), + temperature=0, + max_completion_tokens=100, + model=model_name, + tools=[WEATHER_TOOL], + tool_choice="none", + logprobs=False, + seed=SEED, + stream=True, + ) + + # Pre-v11 models lack grammar enforcement, so the model may still + # emit tool calls even with tool_choice="none". + pre_v11 = _is_pre_v11(server_config) + result = await _collect_streamed_content(stream, no_tool_calls=not pre_v11) + + assert result.finish_reason_count == 1 + if not pre_v11: + assert result.finish_reason != "tool_calls" + streamed_content = "".join(result.chunks) + if not pre_v11: + assert "[TOOL_CALLS]" not in streamed_content + assert streamed_content == non_streaming_content + + +@pytest.mark.asyncio +async def test_chat_without_tools( + client: openai.AsyncOpenAI, server_config: ServerConfig +) -> None: + models = await client.models.list() + model_name: str = models.data[0].id + + # --- non-streaming --- + chat_completion = await client.chat.completions.create( + messages=ensure_system_prompt(MESSAGES_WITHOUT_TOOLS, server_config), + temperature=0, + max_completion_tokens=150, + model=model_name, + logprobs=False, + seed=SEED, + ) + + choice = chat_completion.choices[0] + output_text = choice.message.content + + assert output_text is not None and len(output_text) > 0 + assert choice.finish_reason != "tool_calls" + assert choice.message.tool_calls is None or len(choice.message.tool_calls) == 0 + + # --- streaming --- + stream = await client.chat.completions.create( + messages=ensure_system_prompt(MESSAGES_WITHOUT_TOOLS, server_config), + temperature=0, + max_completion_tokens=150, + model=model_name, + logprobs=False, + seed=SEED, + stream=True, + ) + + result = await _collect_streamed_content( + stream, expected_finish_reason=choice.finish_reason + ) + + assert result.role_sent + assert result.finish_reason_count == 1 + assert len(result.chunks) + assert "".join(result.chunks) == output_text + + +@pytest.mark.asyncio +async def test_tool_call_with_results( + client: openai.AsyncOpenAI, server_config: ServerConfig +) -> None: + _requires_tool_parser(server_config) + + models = await client.models.list() + model_name: str = models.data[0].id + + # --- non-streaming --- + chat_completion = await client.chat.completions.create( + messages=ensure_system_prompt(MESSAGES_WITH_TOOL_RESPONSE, server_config), + temperature=0, + max_completion_tokens=100, + model=model_name, + tools=[WEATHER_TOOL, SEARCH_TOOL], + logprobs=False, + seed=SEED, + ) + + choice = chat_completion.choices[0] + + assert choice.finish_reason != "tool_calls" + assert choice.message.tool_calls is None or len(choice.message.tool_calls) == 0 + assert choice.message.content is not None + assert "98" in choice.message.content + + # --- streaming --- + stream = await client.chat.completions.create( + messages=ensure_system_prompt(MESSAGES_WITH_TOOL_RESPONSE, server_config), + temperature=0, + max_completion_tokens=100, + model=model_name, + tools=[WEATHER_TOOL, SEARCH_TOOL], + logprobs=False, + seed=SEED, + stream=True, + ) + + result = await _collect_streamed_content( + stream, expected_finish_reason=choice.finish_reason + ) + + assert result.role_sent + assert result.finish_reason_count == 1 + assert len(result.chunks) + assert "".join(result.chunks) == choice.message.content + + +def _requires_parallel(server_config: ServerConfig) -> None: + r"""Skip test if the model does not support parallel tool calls.""" + if not server_config.get("supports_parallel"): + pytest.skip( + f"Skipping: {server_config['model']} does not support parallel tool calls" + ) + + +@pytest.mark.asyncio +async def test_tool_call_parallel( + client: openai.AsyncOpenAI, server_config: ServerConfig +) -> None: + _requires_tool_parser(server_config) + _requires_parallel(server_config) + + models = await client.models.list() + model_name: str = models.data[0].id + + # --- non-streaming --- + chat_completion = await client.chat.completions.create( + messages=ensure_system_prompt( + MESSAGES_ASKING_FOR_PARALLEL_TOOLS, server_config + ), + temperature=0, + max_completion_tokens=200, + model=model_name, + tools=[WEATHER_TOOL], + logprobs=False, + seed=SEED, + ) + + choice = chat_completion.choices[0] + tool_calls = choice.message.tool_calls + + assert choice.finish_reason == "tool_calls" + assert tool_calls is not None and len(tool_calls) >= 2 + for tc in tool_calls: + assert tc.type == "function" + assert tc.function.name == "get_current_weather" + assert isinstance(tc.function.arguments, str) + parsed = json.loads(tc.function.arguments) + assert "city" in parsed + assert len(tc.id) == 9 + + non_streaming_tool_calls = tool_calls + + # --- streaming --- + stream = await client.chat.completions.create( + messages=ensure_system_prompt( + MESSAGES_ASKING_FOR_PARALLEL_TOOLS, server_config + ), + temperature=0, + max_completion_tokens=200, + model=model_name, + tools=[WEATHER_TOOL], + logprobs=False, + seed=SEED, + stream=True, + ) + + result = await _collect_streamed_parallel_tool_calls(stream) + + assert result.finish_reason_count == 1 + assert result.role_name == "assistant" + assert len(result.function_names) >= 2 + assert all(name == "get_current_weather" for name in result.function_names) + assert len(result.tool_call_ids) >= 2 + assert all(isinstance(tid, str) and len(tid) == 9 for tid in result.tool_call_ids) + + for args_str in result.function_args_strs: + streamed_args = json.loads(args_str) + assert "city" in streamed_args + + assert len(result.function_names) == len(non_streaming_tool_calls) diff --git a/tests/tool_use/mistral/utils.py b/tests/tool_use/mistral/utils.py index 4d772ba6379..01a2aaee6d2 100644 --- a/tests/tool_use/mistral/utils.py +++ b/tests/tool_use/mistral/utils.py @@ -2,16 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing_extensions import TypedDict - - -class ServerConfig(TypedDict, total=False): - model: str - arguments: list[str] - system_prompt: str | None - supports_parallel: bool | None - supports_rocm: bool | None - +from tests.tool_use.utils import ServerConfig ARGS: list[str] = ["--max-model-len", "1024"] @@ -21,6 +12,11 @@ CONFIGS: dict[str, ServerConfig] = { "arguments": [ "--tokenizer-mode", "mistral", + "--tool-call-parser", + "mistral", + "--enable-auto-tool-choice", + "--enforce-eager", + "--no-enable-prefix-caching", '--ignore-patterns="consolidated.safetensors"', ], "system_prompt": "You are a helpful assistant with access to tools. If a tool" @@ -29,4 +25,22 @@ CONFIGS: dict[str, ServerConfig] = { "without calling a tool. DO NOT CALL A TOOL THAT IS IRRELEVANT " "to the user's question - just respond to it normally.", }, + "ministral-3b": { + "model": "mistralai/Ministral-3-3B-Instruct-2512", + "arguments": [ + "--tokenizer-mode", + "mistral", + "--tool-call-parser", + "mistral", + "--enable-auto-tool-choice", + "--enforce-eager", + "--no-enable-prefix-caching", + ], + "system_prompt": "You are a helpful assistant with access to tools. If a tool" + " that you have would be helpful to answer a user query, " + "call the tool. Otherwise, answer the user's query directly " + "without calling a tool. DO NOT CALL A TOOL THAT IS IRRELEVANT " + "to the user's question - just respond to it normally.", + "supports_parallel": True, + }, } diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 2bc1b6e0875..437e41fb286 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -11,7 +11,7 @@ from openai.types.chat.chat_completion_audio import ( ChatCompletionAudio as OpenAIChatCompletionAudio, ) from openai.types.chat.chat_completion_message import Annotation as OpenAIAnnotation -from pydantic import Field, model_validator +from pydantic import Field, PrivateAttr, model_validator from vllm.config import ModelConfig from vllm.config.utils import replace @@ -398,6 +398,9 @@ class ChatCompletionRequest(OpenAIBaseModel): msg["tool_calls"] = list(tool_calls) return self + _grammar_from_tool_parser: bool = PrivateAttr(default=False) + """CAUTION: Should only be set by ``ToolParser.adjust_request``.""" + def build_chat_params( self, default_template: str | None, @@ -822,13 +825,6 @@ class ChatCompletionRequest(OpenAIBaseModel): return data - @model_validator(mode="before") - @classmethod - def set_include_reasoning_for_none_effort(cls, data: Any) -> Any: - if data.get("reasoning_effort") == "none": - data["include_reasoning"] = False - return data - class BatchChatCompletionRequest(OpenAIBaseModel): """Request model for the /v1/chat/completions/batch endpoint. diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 0b8dd0aa28e..446f127a91e 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -73,7 +73,10 @@ from vllm.reasoning import ReasoningParser from vllm.renderers import ChatParams from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.mistral_tool_parser import MistralToolCall +from vllm.tool_parsers.mistral_tool_parser import ( + MistralToolCall, + MistralToolParser, +) from vllm.tool_parsers.utils import partial_json_loads from vllm.utils.collection_utils import as_list from vllm.utils.mistral import is_mistral_tokenizer @@ -140,6 +143,12 @@ class OpenAIServingChat(OpenAIServing): enable_auto_tools=enable_auto_tools, model_name=self.model_config.model, ) + _is_mistral_tool_parser = self.tool_parser is not None and issubclass( + self.tool_parser, MistralToolParser + ) + if _is_mistral_tool_parser and self.reasoning_parser_cls is not None: + MistralToolParser.model_can_reason = True + self.exclude_tools_when_tool_choice_none = exclude_tools_when_tool_choice_none self.enable_prompt_tokens_details = enable_prompt_tokens_details @@ -310,6 +319,11 @@ class OpenAIServingChat(OpenAIServing): else: if not request.include_reasoning: reasoning_ended = True + elif request._grammar_from_tool_parser: + # The Mistral grammar already includes an optional + # `think?` rule that handles both reasoning and + # non-reasoning outputs. + reasoning_ended = True elif reasoning_parser: reasoning_ended = reasoning_parser.is_reasoning_end( prompt_token_ids or [] @@ -530,6 +544,8 @@ class OpenAIServingChat(OpenAIServing): harmony_tools_streamed = [False] * num_choices tools_streamed = [False] * num_choices + is_mistral_grammar_path = request._grammar_from_tool_parser + if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam): tool_choice_function_name = request.tool_choice.function.name else: @@ -553,7 +569,7 @@ class OpenAIServingChat(OpenAIServing): # Only one of these will be used, thus previous_texts and # all_previous_token_ids will not be used twice in the same iteration. - if tool_choice_auto or reasoning_parser: + if is_mistral_grammar_path or tool_choice_auto or reasoning_parser: # These are only required in "auto" tool choice case all_previous_token_ids = [[] for _ in range(num_choices)] reasoning_end_arr = [False] * num_choices @@ -748,7 +764,7 @@ class OpenAIServingChat(OpenAIServing): delta_message: DeltaMessage | None # just update previous_texts and previous_token_ids - if tool_choice_auto or reasoning_parser: + if is_mistral_grammar_path or tool_choice_auto or reasoning_parser: assert previous_texts is not None assert all_previous_token_ids is not None previous_text = previous_texts[i] @@ -772,6 +788,30 @@ class OpenAIServingChat(OpenAIServing): ) ) harmony_tools_streamed[i] |= tools_streamed_flag + # Mistral grammar path: combined reasoning + tool streaming + elif is_mistral_grammar_path: + assert tool_parser is not None + assert isinstance(tool_parser, MistralToolParser) + assert reasoning_end_arr is not None + output_token_ids = as_list(output.token_ids) + result = tool_parser.extract_maybe_reasoning_and_tool_streaming( + reasoning_parser=reasoning_parser, + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=previous_token_ids, + current_token_ids=current_token_ids, + output_token_ids=output_token_ids, + reasoning_ended=reasoning_end_arr[i], + prompt_is_reasoning_end=(prompt_is_reasoning_end_arr[i]), + request=request, + ) + delta_message = result.delta_message + reasoning_end_arr[i] = result.reasoning_ended + current_text = result.current_text + current_token_ids = result.current_token_ids + if result.tools_called: + tools_streamed[i] = True # handle streaming deltas for tools with named tool_choice elif tool_choice_function_name: # When encountering think end id in prompt_token_ids @@ -925,7 +965,9 @@ class OpenAIServingChat(OpenAIServing): delta_message = DeltaMessage(content=delta_text) # update the previous values for the next iteration - if (tool_choice_auto or reasoning_parser) and not self.use_harmony: + if ( + is_mistral_grammar_path or tool_choice_auto or reasoning_parser + ) and not self.use_harmony: assert previous_texts is not None assert all_previous_token_ids is not None previous_texts[i] = current_text @@ -1312,7 +1354,24 @@ class OpenAIServingChat(OpenAIServing): tool_call_class = ( MistralToolCall if is_mistral_tokenizer(tokenizer) else ToolCall ) - if (not self.enable_auto_tools or not self.tool_parser) and ( + + use_mistral_tool_parser = request._grammar_from_tool_parser + if use_mistral_tool_parser: + tool_call_items = MistralToolParser.build_non_streaming_tool_calls( + tool_calls + ) + if tool_call_items: + auto_tools_called = ( + request.tool_choice is None or request.tool_choice == "auto" + ) + message = ChatMessage( + role=role, + reasoning=reasoning, + content=content, + tool_calls=tool_call_items, + ) + + elif (not self.enable_auto_tools or not self.tool_parser) and ( not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) and request.tool_choice != "required" ): diff --git a/vllm/entrypoints/openai/engine/serving.py b/vllm/entrypoints/openai/engine/serving.py index 5bd415b4fbd..85237604e5a 100644 --- a/vllm/entrypoints/openai/engine/serving.py +++ b/vllm/entrypoints/openai/engine/serving.py @@ -65,6 +65,7 @@ from vllm.renderers.inputs.preprocess import ( from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike from vllm.tool_parsers import ToolParser +from vllm.tool_parsers.mistral_tool_parser import MistralToolParser from vllm.tracing import ( contains_trace_headers, extract_trace_headers, @@ -610,16 +611,31 @@ class OpenAIServing: tool_parser_cls: type[ToolParser] | None, content: str | None = None, ) -> tuple[list[FunctionCall] | None, str | None]: + # When the Mistral grammar factory injected structured outputs, + # let the parser handle the output. + use_mistral_tool_parser = ( + isinstance(request, ChatCompletionRequest) + and tool_parser_cls is not None + and issubclass(tool_parser_cls, MistralToolParser) + and request._grammar_from_tool_parser + ) + function_calls = list[FunctionCall]() - if request.tool_choice and isinstance(request.tool_choice, ToolChoiceFunction): + if ( + not use_mistral_tool_parser + and request.tool_choice + and isinstance(request.tool_choice, ToolChoiceFunction) + ): assert content is not None # Forced Function Call function_calls.append( FunctionCall(name=request.tool_choice.name, arguments=content) ) content = None # Clear content since tool is called. - elif request.tool_choice and isinstance( - request.tool_choice, ChatCompletionNamedToolChoiceParam + elif ( + not use_mistral_tool_parser + and request.tool_choice + and isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) ): assert content is not None # Forced Function Call @@ -627,7 +643,7 @@ class OpenAIServing: FunctionCall(name=request.tool_choice.function.name, arguments=content) ) content = None # Clear content since tool is called. - elif request.tool_choice == "required": + elif not use_mistral_tool_parser and request.tool_choice == "required": tool_calls = [] with contextlib.suppress(ValidationError): content = content or "" @@ -642,10 +658,12 @@ class OpenAIServing: ) ) content = None # Clear content since tool is called. - elif ( - tool_parser_cls - and enable_auto_tools - and (request.tool_choice == "auto" or request.tool_choice is None) + elif tool_parser_cls and ( + use_mistral_tool_parser + or ( + enable_auto_tools + and (request.tool_choice == "auto" or request.tool_choice is None) + ) ): if tokenizer is None: raise ValueError( diff --git a/vllm/entrypoints/serve/render/serving.py b/vllm/entrypoints/serve/render/serving.py index 5aa2449797b..99b6cb47cb5 100644 --- a/vllm/entrypoints/serve/render/serving.py +++ b/vllm/entrypoints/serve/render/serving.py @@ -53,6 +53,7 @@ from vllm.renderers.inputs.preprocess import ( prompt_to_seq, ) from vllm.tool_parsers import ToolParser +from vllm.tool_parsers.mistral_tool_parser import MistralToolParser from vllm.utils import random_uuid from vllm.utils.mistral import is_mistral_tokenizer from vllm.utils.mistral import mt as _mt @@ -555,9 +556,19 @@ class OpenAIServingRender: # tool parsing is done only if a tool_parser has been set and if # tool_choice is not "none" (if tool_choice is "none" but a tool_parser # is set, we want to prevent parsing a tool_call hallucinated by the LLM + # + # Exception: Mistral grammar-capable tokenizers always call + # adjust_request — even for tool_choice="none" — so that the grammar + # factory can prevent special-token leakage. if tool_parser is not None: tool_choice = getattr(request, "tool_choice", "none") - if tool_choice != "none": + tokenizer = renderer.get_tokenizer() + is_mistral_grammar_eligible = ( + issubclass(tool_parser, MistralToolParser) + and is_mistral_tokenizer(tokenizer) + and tokenizer.supports_grammar + ) + if tool_choice != "none" or is_mistral_grammar_eligible: if not isinstance(request, ChatCompletionRequest | ResponsesRequest): msg = ( "Tool usage is only supported " @@ -565,7 +576,6 @@ class OpenAIServingRender: f"but got {type(request).__name__}" ) raise NotImplementedError(msg) - tokenizer = renderer.get_tokenizer() request = tool_parser(tokenizer, request.tools).adjust_request( request=request ) diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 9bcc669591e..77fa6402180 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -157,6 +157,10 @@ def _is_non_tekken_mistral(tokenizer: TokenizerLike) -> bool: return is_mistral_tokenizer(tokenizer) and not tokenizer.is_tekken +def _get_llg_tokenizer(tokenizer: TokenizerLike) -> Any: + return tokenizer.llg_tokenizer if is_mistral_tokenizer(tokenizer) else None + + class SamplingParams( PydanticMsgspecMixin, msgspec.Struct, @@ -816,7 +820,10 @@ class SamplingParams( # allows <|special_token|> and similar, see # https://github.com/guidance-ai/llguidance/blob/main/docs/syntax.md#special-tokens # Without tokenizer these are disallowed in grammars. - validate_guidance_grammar(self, tokenizer=None) + validate_guidance_grammar( + self, + tokenizer=_get_llg_tokenizer(tokenizer), + ) elif backend == "outlines": # outlines backend validate_structured_output_request_outlines(self) @@ -862,7 +869,10 @@ class SamplingParams( self.structured_outputs._backend = "outlines" else: # Fall back to guidance by default. - validate_guidance_grammar(self, tokenizer=None) + validate_guidance_grammar( + self, + tokenizer=_get_llg_tokenizer(tokenizer), + ) self.structured_outputs._backend = "guidance" # Remember that this backend was set automatically self.structured_outputs._backend_was_auto = True diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py index 147dca88877..ef58b1b75d6 100644 --- a/vllm/tokenizers/mistral.py +++ b/vllm/tokenizers/mistral.py @@ -54,6 +54,50 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _pop_unallowed_keys_and_warn( + dictionary: dict[str, Any], allowed_keys: set[str], err_dict_name: str +): + keys = list(dictionary.keys()) + for key in keys: + if key not in allowed_keys: + dictionary.pop(key) + logger.warning_once( + f"'{key=}' is not supported by mistral-common " + f"for {err_dict_name}. It has been popped from the " + "object." + ) + + +# TODO(juliendenize): remove this once OpenAI API is better supported by +# `mistral-common`. +def adapt_inplace_to_mistral_tool( + tool: dict[str, Any], +) -> dict[str, Any]: + tools_fields = set(Tool.model_fields.keys()) + function_fields = set(Function.model_fields.keys()) + + # The Mistral client, in comparison to the OpenAI client, requires the + # "parameters" dict and the "description" string to be present + # even if they are empty. + if function := tool.get("function"): + if function.get("parameters") is None: + function["parameters"] = {} + if function.get("description") is None: + function["description"] = "" + + _pop_unallowed_keys_and_warn( + dictionary=function, + allowed_keys=function_fields, + err_dict_name="function", + ) + + _pop_unallowed_keys_and_warn( + dictionary=tool, allowed_keys=tools_fields, err_dict_name="tools" + ) + + return tool + + def maybe_serialize_tool_calls(request: "MistralChatCompletionRequest"): # SEE: https://github.com/vllm-project/vllm/pull/9951 # Credits go to: @gcalmettes @@ -159,44 +203,11 @@ def _prepare_apply_chat_template_tools_and_messages( # Remove reasoning as unsupported by Mistral _ = message.pop("reasoning", None) # type: ignore - # The Mistral client, in comparison to the OpenAI client, requires the - # "parameters" dict and the "description" string to be present - # even if they are empty. - if tools: - for function in [ - tool["function"] for tool in tools if tool["type"] == "function" - ]: - if function.get("parameters") is None: - function["parameters"] = {} - if function.get("description") is None: - function["description"] = "" - - # We filter not supported arguments to avoid throwing an error. - # TODO(juliendenize): remove this once OpenAI API is better supported by - # `mistral-common`. - tools_fields = set(Tool.model_fields.keys()) - function_fields = set(Function.model_fields.keys()) - for tool in tools: - tool_keys = list(tool.keys()) - for tool_key in tool_keys: - if tool_key not in tools_fields: - tool.pop(tool_key) - logger.warning_once( - f"'{tool_key}' is not supported by mistral-common for tools. " - "It has been popped from the tool definition." - ) - if tool["type"] == "function": - function_keys = list(tool["function"].keys()) - for function_key in function_keys: - if function_key not in function_fields: - tool["function"].pop(function_key) - logger.warning_once( - f"'{function_key}' is not supported by mistral-common " - "for function tools. It has been popped from the " - "function definition." - ) - else: - raise ValueError("mistral-common only supports function tools.") + tools = ( + [adapt_inplace_to_mistral_tool(tool=tool) for tool in tools] + if tools is not None + else None + ) return messages, tools diff --git a/vllm/tool_parsers/mistral_tool_parser.py b/vllm/tool_parsers/mistral_tool_parser.py index 4d1aaffedd0..1d2613104fc 100644 --- a/vllm/tool_parsers/mistral_tool_parser.py +++ b/vllm/tool_parsers/mistral_tool_parser.py @@ -1,12 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + import json from collections.abc import Sequence +from dataclasses import dataclass from enum import Enum, auto from random import choices from string import ascii_letters, digits -from typing import Any +from typing import TYPE_CHECKING, Any import ijson import regex as re @@ -37,14 +40,19 @@ from vllm.entrypoints.openai.engine.protocol import ( ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.reasoning.mistral_reasoning_parser import MistralReasoningParser from vllm.sampling_params import StructuredOutputsParams from vllm.tokenizers import TokenizerLike +from vllm.tokenizers.mistral import MistralTokenizer, adapt_inplace_to_mistral_tool from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, ) from vllm.utils.mistral import is_mistral_tokenizer +if TYPE_CHECKING: + from vllm.reasoning import ReasoningParser + logger = init_logger(__name__) ALPHANUMERIC = ascii_letters + digits @@ -86,13 +94,28 @@ def _is_pre_v11_tokeniser(model_tokenizer: TokenizerLike) -> bool: return not (is_mistral_tokenizer(model_tokenizer) and model_tokenizer.version >= 11) -class MistralToolParser(ToolParser): +@dataclass +class MistralStreamingResult: + r"""Encapsulates the mutable state returned from + `MistralToolParser.extract_maybe_reasoning_and_tool_streaming`. """ - Tool call parser for Mistral 7B Instruct v0.3, intended for use with - - [`mistral_common`](https://github.com/mistralai/mistral-common/) - - the examples/tool_chat_template_mistral.jinja template. - Used when --enable-auto-tool-choice --tool-call-parser mistral are all set + delta_message: DeltaMessage | None + reasoning_ended: bool + tools_called: bool + current_text: str + current_token_ids: list[int] + + +class MistralToolParser(ToolParser): + r"""Tool call parser for Mistral models, intended for use with either: + + - `mistral_common `_ + (recommended) + - the `examples/tool_chat_template_mistral.jinja` template. + + Used when `--enable-auto-tool-choice --tool-call-parser mistral` are all + set. """ # Used to generate correct grammar in `adjust_request` @@ -210,9 +233,11 @@ class MistralToolParser(ToolParser): reasoning=self.model_can_reason ) - tools = ( + mistral_tools = ( [ - MistralTool.from_openai(openai_tool=tool.model_dump()) + MistralTool.model_validate( + adapt_inplace_to_mistral_tool(tool.model_dump()) + ) for tool in request.tools ] if request.tools is not None @@ -244,15 +269,158 @@ class MistralToolParser(ToolParser): lark_grammar = grammar_factory.get_lark_from_jinja( template=template, mode=tool_choice, - tools=tools, + tools=mistral_tools, json_schema=json_schema, parallel_tool_calls=request.parallel_tool_calls, json_only=False, ) request.structured_outputs = StructuredOutputsParams(grammar=lark_grammar) + request._grammar_from_tool_parser = True return request + def extract_maybe_reasoning_and_tool_streaming( + self, + *, + reasoning_parser: ReasoningParser | None, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: list[int], + current_token_ids: list[int], + output_token_ids: Sequence[int], + reasoning_ended: bool, + prompt_is_reasoning_end: bool | None, + request: ChatCompletionRequest, + ) -> MistralStreamingResult: + r"""Streaming extraction with reasoning followed by tool-call parsing. + + This method encapsulates the combined reasoning extraction and + tool-call streaming logic so that the serving layer only needs a + thin routing branch. + + The flow is: + + 1. If a *reasoning_parser* is present and reasoning has **not** ended, + extract reasoning tokens. Pre-v15 models may have pre-filled + `[THINK]...[/THINK]` in system prompts, so we skip the + prompt-level reasoning-end check for those. + 2. Once reasoning ends (or if there is no reasoning parser), delegate + to `extract_tool_calls_streaming` and track whether tools were + called. + + Args: + reasoning_parser: Optional reasoning parser instance. + previous_text: Accumulated text from prior chunks. + current_text: Full accumulated text including current chunk. + delta_text: New text in this chunk. + previous_token_ids: Token ids from prior chunks. + current_token_ids: Full token ids including current chunk. + output_token_ids: Raw output token ids from the engine. + reasoning_ended: Whether reasoning has already ended. + prompt_is_reasoning_end: Whether the prompt itself ends reasoning. + request: The originating chat completion request. + """ + delta_message: DeltaMessage | None = None + tools_called = False + reasoning_ended_at_entry = reasoning_ended + + # For MistralReasoningParser, only enter the reasoning block when + # the model has actually emitted a [THINK] token. Other reasoning + # parsers always expect thinking to be present. + expect_thinking = ( + not isinstance(reasoning_parser, MistralReasoningParser) + or reasoning_parser.start_token_id in current_token_ids + ) + if reasoning_parser is not None and not reasoning_ended and expect_thinking: + # Pre-v15 models may have pre-filled [THINK]...[/THINK] in + # system prompts, so skip the prompt-level reasoning-end + # check and wait for the output's own end-of-think. + is_pre_v15 = ( + isinstance(self.model_tokenizer, MistralTokenizer) + and self.model_tokenizer.version < 15 + ) + + if not is_pre_v15 and prompt_is_reasoning_end: + reasoning_ended = True + current_token_ids = list(output_token_ids) + else: + delta_message = reasoning_parser.extract_reasoning_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + output_token_ids, + ) + if reasoning_parser.is_reasoning_end_streaming( + current_token_ids, output_token_ids + ): + reasoning_ended = True + current_token_ids = reasoning_parser.extract_content_ids( + list(output_token_ids) + ) + if delta_message and delta_message.content: + current_text = delta_message.content + delta_message.content = None + else: + current_text = "" + + if not reasoning_ended: + return MistralStreamingResult( + delta_message=delta_message, + reasoning_ended=False, + tools_called=False, + current_text=current_text, + current_token_ids=current_token_ids, + ) + + delta_token_ids = list(output_token_ids) + + # On the iteration where reasoning just ended, reset the text/token + # state so the tool parser sees a clean history instead of the + # accumulated reasoning text. + if not reasoning_ended_at_entry and reasoning_ended: + previous_text = "" + previous_token_ids = [] + delta_text = current_text + delta_token_ids = current_token_ids + + delta_message = self.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=previous_token_ids, + current_token_ids=current_token_ids, + delta_token_ids=delta_token_ids, + request=request, + ) + if delta_message and delta_message.tool_calls: + tools_called = True + + return MistralStreamingResult( + delta_message=delta_message, + reasoning_ended=reasoning_ended, + tools_called=tools_called, + current_text=current_text, + current_token_ids=current_token_ids, + ) + + @staticmethod + def build_non_streaming_tool_calls( + tool_calls: list[FunctionCall] | None, + ) -> list[ToolCall]: + r"""Build `MistralToolCall` items for non-streaming responses.""" + if not tool_calls: + return [] + + return [ + MistralToolCall(id=tc.id, function=tc) + if tc.id + else MistralToolCall(function=tc) + for tc in tool_calls + ] + def extract_tool_calls( self, model_output: str, @@ -323,7 +491,7 @@ class MistralToolParser(ToolParser): )[0] tool_calls = json.loads(raw_tool_call) except (IndexError, json.JSONDecodeError): - logger.exception("Error in extracting tool call from response: {e}") + logger.exception("Error in extracting tool call from response.") # If raw decoding and decoding post regex rule fails, then just # return content. return ExtractedToolCallInformation( From 18013df6ae27c3fb941307c46c975227126d641f Mon Sep 17 00:00:00 2001 From: jigangz <115519042+jigangz@users.noreply.github.com> Date: Wed, 15 Apr 2026 21:08:04 -0700 Subject: [PATCH 041/150] [Bugfix] Reject empty tools array with HTTP 400 (#39780) Signed-off-by: Jigang Zhou Co-authored-by: Claude --- .../test_ernie45_moe_tool_parser.py | 2 +- tests/tool_parsers/test_xlam_tool_parser.py | 2 +- ...est_chat_completion_request_validations.py | 18 +++++++------- .../openai/chat_completion/protocol.py | 24 +++++++++---------- 4 files changed, 23 insertions(+), 23 deletions(-) diff --git a/tests/tool_parsers/test_ernie45_moe_tool_parser.py b/tests/tool_parsers/test_ernie45_moe_tool_parser.py index a00e4389476..ffee62441c9 100644 --- a/tests/tool_parsers/test_ernie45_moe_tool_parser.py +++ b/tests/tool_parsers/test_ernie45_moe_tool_parser.py @@ -328,7 +328,7 @@ def test_extract_tool_calls_streaming_incremental( expected_content, ): """Verify the Ernie45 Parser streaming behavior by verifying each chunk is as expected.""" # noqa: E501 - request = ChatCompletionRequest(model=MODEL, messages=[], tools=[]) + request = ChatCompletionRequest(model=MODEL, messages=[]) tool_calls_dict = {} for delta_message in stream_delta_message_generator( diff --git a/tests/tool_parsers/test_xlam_tool_parser.py b/tests/tool_parsers/test_xlam_tool_parser.py index a5cab218f72..3853d2039a7 100644 --- a/tests/tool_parsers/test_xlam_tool_parser.py +++ b/tests/tool_parsers/test_xlam_tool_parser.py @@ -484,7 +484,7 @@ def test_extract_tool_calls_streaming_incremental( expected_content, ): """Verify the XLAM Parser streaming behavior by verifying each chunk is as expected.""" # noqa: E501 - request = ChatCompletionRequest(model=MODEL, messages=[], tools=[]) + request = ChatCompletionRequest(model=MODEL, messages=[]) chunks = [] for delta_message in stream_delta_message_generator( diff --git a/tests/tool_use/test_chat_completion_request_validations.py b/tests/tool_use/test_chat_completion_request_validations.py index 69846f9adb1..70f4ac22541 100644 --- a/tests/tool_use/test_chat_completion_request_validations.py +++ b/tests/tool_use/test_chat_completion_request_validations.py @@ -26,15 +26,15 @@ def test_chat_completion_request_with_no_tools(): ) assert request.tool_choice == "none" - # tools key present but empty - request = ChatCompletionRequest.model_validate( - { - "messages": [{"role": "user", "content": "Hello"}], - "model": "facebook/opt-125m", - "tools": [], - } - ) - assert request.tool_choice == "none" + # tools key present but empty -- should be rejected + with pytest.raises(ValueError, match="must not be an empty array"): + ChatCompletionRequest.model_validate( + { + "messages": [{"role": "user", "content": "Hello"}], + "model": "facebook/opt-125m", + "tools": [], + } + ) @pytest.mark.parametrize("tool_choice", ["auto", "required"]) diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 437e41fb286..aacac38e07f 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -681,6 +681,18 @@ class ChatCompletionRequest(OpenAIBaseModel): @model_validator(mode="before") @classmethod def check_tool_usage(cls, data): + if isinstance(data, ValueError): + raise data + if not isinstance(data, dict): + return data + + # Reject empty tools array, matching OpenAI API behavior + if data.get("tools") == []: + raise ValueError( + "`tools` must not be an empty array. " + "Either provide at least one tool or omit the field entirely." + ) + # if "tool_choice" is not specified but tools are provided, # default to "auto" tool_choice if "tool_choice" not in data and data.get("tools"): @@ -707,18 +719,6 @@ class ChatCompletionRequest(OpenAIBaseModel): "are supported." ) - # if tool_choice is "required" but the "tools" list is empty, - # override the data to behave like "none" to align with - # OpenAI’s behavior. - if ( - data["tool_choice"] == "required" - and isinstance(data["tools"], list) - and len(data["tools"]) == 0 - ): - data["tool_choice"] = "none" - del data["tools"] - return data - # ensure that if "tool_choice" is specified as an object, # it matches a valid tool correct_usage_message = ( From 445b7093fdbbc6c52875d2deef28df1d5da56089 Mon Sep 17 00:00:00 2001 From: Fadi Arafeh <115173828+fadara01@users.noreply.github.com> Date: Thu, 16 Apr 2026 07:26:17 +0200 Subject: [PATCH 042/150] [perf][cpu] Accelerate BF16 GELU with LUT impl on Arm CPUs (#37469) Signed-off-by: Fadi Arafeh Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../scripts/hardware_ci/run-cpu-test-arm.sh | 1 + cmake/cpu_extension.cmake | 1 + csrc/cpu/activation_lut_bf16.cpp | 71 +++++++++++ csrc/cpu/torch_bindings.cpp | 12 ++ csrc/ops.h | 1 + tests/kernels/core/test_cpu_activation.py | 111 ++++++++++++++++++ vllm/_custom_ops.py | 6 + vllm/model_executor/layers/activation.py | 32 ++++- vllm/platforms/cpu.py | 7 ++ 9 files changed, 240 insertions(+), 2 deletions(-) create mode 100644 csrc/cpu/activation_lut_bf16.cpp create mode 100644 tests/kernels/core/test_cpu_activation.py diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index c2509a07b2c..7166435ac1e 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -51,6 +51,7 @@ function cpu_tests() { set -e pytest -x -v -s tests/kernels/test_onednn.py pytest -x -v -s tests/kernels/attention/test_cpu_attn.py + pytest -x -v -s tests/kernels/core/test_cpu_activation.py pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic" # basic online serving diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 0389ce1299a..f45280ee488 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -360,6 +360,7 @@ set(VLLM_EXT_SRC if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) set(VLLM_EXT_SRC "csrc/cpu/shm.cpp" + "csrc/cpu/activation_lut_bf16.cpp" ${VLLM_EXT_SRC}) endif() diff --git a/csrc/cpu/activation_lut_bf16.cpp b/csrc/cpu/activation_lut_bf16.cpp new file mode 100644 index 00000000000..0ff2567e1ee --- /dev/null +++ b/csrc/cpu/activation_lut_bf16.cpp @@ -0,0 +1,71 @@ +#include "cpu_types.hpp" + +#include +#include +#include +#include + +#include +#include +#include + +constexpr uint32_t ActivationLutSize = 1u << 16; + +at::Tensor gelu_reference(const at::Tensor& x) { return at::gelu(x, "none"); } + +void maybe_init_activation_lut_bf16( + uint16_t* lut, std::once_flag& once, + at::Tensor (*activation)(const at::Tensor&)) { + std::call_once(once, [&]() { + auto lut_input = + at::empty({static_cast(ActivationLutSize)}, + at::TensorOptions().device(at::kCPU).dtype(at::kFloat)); + auto* lut_input_ptr = lut_input.data_ptr(); +#pragma omp parallel for + for (uint32_t i = 0; i < ActivationLutSize; ++i) { + lut_input_ptr[i] = c10::detail::f32_from_bits(static_cast(i)); + } + + auto lut_output = activation(lut_input); + const auto* lut_output_ptr = lut_output.data_ptr(); +#pragma omp parallel for + for (uint32_t i = 0; i < ActivationLutSize; ++i) { + lut[i] = c10::detail::round_to_nearest_even(lut_output_ptr[i]); + } + }); +} + +void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input, + const uint16_t* lut, const char* op_name) { + TORCH_CHECK(input.scalar_type() == at::kBFloat16, op_name, + ": input must be bfloat16"); + TORCH_CHECK(out.scalar_type() == at::kBFloat16, op_name, + ": out must be bfloat16"); + TORCH_CHECK(input.is_contiguous(), op_name, ": input must be contiguous"); + TORCH_CHECK(out.is_contiguous(), op_name, ": out must be contiguous"); + + const auto* src = + reinterpret_cast(input.data_ptr()); + auto* dst = reinterpret_cast(out.data_ptr()); + const int64_t n = input.numel(); + + CPU_KERNEL_GUARD_IN(activation_lut_bf16_impl) +#pragma omp parallel for + for (int64_t i = 0; i < n; ++i) { + dst[i] = lut[src[i]]; + } + CPU_KERNEL_GUARD_OUT(activation_lut_bf16_impl) +} + +void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input, + const std::string& activation) { + if (activation == "gelu") { + static std::array lut{}; + static std::once_flag once; + maybe_init_activation_lut_bf16(lut.data(), once, gelu_reference); + activation_lut_bf16(out, input, lut.data(), "gelu_lut"); + return; + } + + TORCH_CHECK(false, "Unsupported activation: ", activation); +} diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index fcf7064f606..83bb57c9c09 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -85,6 +85,9 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional bias); +void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input, + const std::string& activation); + torch::Tensor get_scheduler_metadata( const int64_t num_req, const int64_t num_heads_q, const int64_t num_heads_kv, const int64_t head_dim, @@ -231,6 +234,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); ops.impl("gelu_quick", torch::kCPU, &gelu_quick); +#if (defined(__aarch64__) && !defined(__APPLE__)) + + ops.def( + "activation_lut_bf16(Tensor! out, Tensor input, str activation)" + " -> ()"); + ops.impl("activation_lut_bf16", torch::kCPU, &activation_lut_bf16); + +#endif // (defined(__aarch64__) && !defined(__APPLE__)) + // Layernorm // Apply Root Mean Square (RMS) Normalization to the input tensor. ops.def( diff --git a/csrc/ops.h b/csrc/ops.h index 0a0b6c2d7d0..f101ab6fd92 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include diff --git a/tests/kernels/core/test_cpu_activation.py b/tests/kernels/core/test_cpu_activation.py new file mode 100644 index 00000000000..40b5f045468 --- /dev/null +++ b/tests/kernels/core/test_cpu_activation.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from tests.kernels.allclose_default import get_default_atol, get_default_rtol +from tests.kernels.utils import opcheck +from vllm.platforms import CpuArchEnum, current_platform +from vllm.utils.torch_utils import set_random_seed + +if not current_platform.is_cpu(): + pytest.skip("skipping CPU-only tests", allow_module_level=True) + +from vllm.model_executor.layers.activation import ( + GELU, + FastGELU, + GeluAndMul, + NewGELU, + QuickGELU, + SiluAndMul, +) + +DTYPES = [torch.bfloat16, torch.float32] +NUM_TOKENS = [7, 83] +D = [512, 2048] +SEEDS = [0] + + +@pytest.mark.parametrize( + ("activation_cls", "fn"), + [ + (SiluAndMul, torch.ops._C.silu_and_mul), + (GeluAndMul, torch.ops._C.gelu_and_mul), + (GeluAndMul, torch.ops._C.gelu_tanh_and_mul), + ], +) +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("d", D) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@torch.inference_mode() +def test_cpu_act_and_mul( + default_vllm_config, + activation_cls: type[torch.nn.Module], + fn: object, + num_tokens: int, + d: int, + dtype: torch.dtype, + seed: int, +) -> None: + set_random_seed(seed) + x = torch.randn(num_tokens, 2 * d, dtype=dtype) + + layer = activation_cls() + out = layer(x) + ref_out = layer.forward_native(x) + + torch.testing.assert_close( + out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out) + ) + + output_shape = x.shape[:-1] + (x.shape[-1] // 2,) + raw_out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + opcheck(fn, (raw_out, x)) + + +@pytest.mark.parametrize( + ("activation_cls", "fn", "op_args"), + [ + (NewGELU, torch.ops._C.gelu_new, ()), + (FastGELU, torch.ops._C.gelu_fast, ()), + (QuickGELU, torch.ops._C.gelu_quick, ()), + pytest.param( + GELU, + getattr(torch.ops._C, "activation_lut_bf16", None), + ("gelu",), + marks=pytest.mark.skipif( + current_platform.get_cpu_architecture() != CpuArchEnum.ARM, + reason="activation_lut_bf16 is only built on Arm CPU", + ), + ), + ], +) +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("d", D) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@torch.inference_mode() +def test_cpu_unary_activation( + default_vllm_config, + activation_cls: type[torch.nn.Module], + fn: object, + op_args: tuple[str, ...], + num_tokens: int, + d: int, + dtype: torch.dtype, + seed: int, +) -> None: + set_random_seed(seed) + x = torch.randn(num_tokens, d, dtype=dtype) + layer = activation_cls() + out = layer(x) + ref_out = layer.forward_native(x) + torch.testing.assert_close( + out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out) + ) + # gelu with activation_lut_bf16 only makes sense for BF16 + if not (activation_cls is GELU and dtype != torch.bfloat16): + raw_out = torch.empty_like(x) + opcheck(fn, (raw_out, x, *op_args)) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index d6780185be9..facac5a192d 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3300,6 +3300,12 @@ def cpu_gemm_wna16( return output +def cpu_activation_lut_bf16(input: torch.Tensor, activation: str) -> torch.Tensor: + out = torch.empty_like(input) + torch.ops._C.activation_lut_bf16(out, input, activation) + return out + + def cpu_prepack_moe_weight( weight: torch.Tensor, isa: str, diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index 3e00d21d5a1..26a771cb750 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -16,7 +16,7 @@ from vllm.distributed import ( from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.utils import set_weight_attrs -from vllm.platforms import current_platform +from vllm.platforms import CpuArchEnum, current_platform from vllm.triton_utils import tl, triton from vllm.utils.collection_utils import LazyDict @@ -247,6 +247,34 @@ class GeluAndMulSparse(CustomOp): return self.forward_native(x) +# --8<-- [start:gelu] +@CustomOp.register("gelu") +class GELU(CustomOp): + # --8<-- [end:gelu] + + def __init__(self): + super().__init__() + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM and hasattr( + torch.ops._C, "activation_lut_bf16" + ): + self.op = torch.ops._C.activation_lut_bf16 + else: + self.op = None + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + return F.gelu(x, approximate="none") + + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if self.op and x.dtype == torch.bfloat16 and x.is_contiguous(): + out = torch.empty_like(x) + self.op(out, x, "gelu") + return out + return self.forward_native(x) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + # --8<-- [start:gelu_and_mul] @CustomOp.register("gelu_and_mul") class GeluAndMul(CustomOp): @@ -635,7 +663,7 @@ class ScaledActivation(nn.Module): _ACTIVATION_REGISTRY = LazyDict( { - "gelu": lambda: nn.GELU(), + "gelu": lambda: GELU(), "gelu_fast": lambda: FastGELU(), "gelu_new": lambda: NewGELU(), "gelu_pytorch_tanh": lambda: ( diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index eb4a5e14827..f319dbc497a 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -246,6 +246,13 @@ class CpuPlatform(Platform): if vllm_config.lora_config is not None: compilation_config.mode = CompilationMode.NONE + if ( + cls.get_cpu_architecture() == CpuArchEnum.ARM + and "+gelu" not in compilation_config.custom_ops + and "-gelu" not in compilation_config.custom_ops + ): + compilation_config.custom_ops.append("+gelu") + vllm_config.profiler_config.torch_profiler_dump_cuda_time_total = False assert vllm_config.device_config.device_type == "cpu" From 4b7ca37bd45467c6190dc139e58554b389d80ac1 Mon Sep 17 00:00:00 2001 From: R3hankhan Date: Thu, 16 Apr 2026 10:56:21 +0530 Subject: [PATCH 043/150] [CPU][IBM Z][Dockefile][Docs] Fix s390x builds for torch 2.11 and update docs for s390x (#39910) Signed-off-by: Rehan Khan --- csrc/cpu/cpu_attn_impl.hpp | 3 + csrc/cpu/utils.hpp | 22 ++++++ docker/Dockerfile.s390x | 68 +++++++++---------- .../installation/cpu.s390x.inc.md | 24 +++---- requirements/common.txt | 2 +- 5 files changed, 72 insertions(+), 47 deletions(-) diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index c15799fa950..08f42459e14 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -147,6 +147,9 @@ struct AttentionMetadata { case ISA::NEON: ss << "NEON, "; break; + case ISA::VXE: + ss << "VXE, "; + break; } ss << "workitem_group_num: " << workitem_group_num << ", reduction_item_num: " << reduction_item_num diff --git a/csrc/cpu/utils.hpp b/csrc/cpu/utils.hpp index f237bba088b..2c9e01c60f9 100644 --- a/csrc/cpu/utils.hpp +++ b/csrc/cpu/utils.hpp @@ -54,12 +54,34 @@ struct Counter { }; inline int64_t get_available_l2_size() { +#if defined(__s390x__) + static int64_t size = []() { + uint32_t l2_cache_size = 0; + auto caps = at::cpu::get_cpu_capabilities(); + auto it = caps.find("l2_cache_size"); + if (it != caps.end()) { + l2_cache_size = static_cast(it->second.toInt()); + } + if (l2_cache_size == 0) { + long sys_l2 = sysconf(_SC_LEVEL2_CACHE_SIZE); + if (sys_l2 > 0) { + l2_cache_size = static_cast(sys_l2); + } + } + if (l2_cache_size == 0) { + l2_cache_size = 256 * 1024; + } + return static_cast(l2_cache_size) >> 1; // use 50% of L2 cache + }(); + return size; +#else static int64_t size = []() { auto caps = at::cpu::get_cpu_capabilities(); const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt(); return l2_cache_size >> 1; // use 50% of L2 cache }(); return size; +#endif } template diff --git a/docker/Dockerfile.s390x b/docker/Dockerfile.s390x index 6b5d965be76..554a7257c23 100644 --- a/docker/Dockerfile.s390x +++ b/docker/Dockerfile.s390x @@ -42,7 +42,7 @@ FROM python-install AS pyarrow # Build Apache Arrow WORKDIR /tmp RUN --mount=type=cache,target=/root/.cache/uv \ - git clone https://github.com/apache/arrow.git && \ + git clone https://github.com/apache/arrow.git -b maint-19.0.1 && \ cd arrow/cpp && \ mkdir release && cd release && \ cmake -DCMAKE_BUILD_TYPE=Release \ @@ -68,19 +68,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install -r requirements-build.txt && \ python setup.py build_ext --build-type=$ARROW_BUILD_TYPE --bundle-arrow-cpp bdist_wheel -FROM python-install AS numa-build -# Install numactl (needed for numa.h dependency) -WORKDIR /tmp -RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.16.tar.gz && \ - tar -xvzf v2.0.16.tar.gz && \ - cd numactl-2.0.16 && \ - ./autogen.sh && \ - ./configure && \ - make - -# Set include path -ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH" - FROM python-install AS rust ENV CARGO_HOME=/root/.cargo ENV RUSTUP_HOME=/root/.rustup @@ -91,6 +78,18 @@ RUN curl https://sh.rustup.rs -sSf | sh -s -- -y && \ rustup default stable && \ rustup show +FROM python-install AS numa-build +WORKDIR /tmp +RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz && \ + tar -xvzf v2.0.19.tar.gz && \ + cd numactl-2.0.19 && \ + ./autogen.sh && \ + ./configure && \ + make + +# Set include path +ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH" + FROM python-install AS torch-vision # Install torchvision ARG TORCH_VISION_VERSION=v0.26.0 @@ -133,7 +132,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ git clone --recursive https://github.com/numba/llvmlite.git -b v0.44.0 && \ git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \ cd llvm-project && mkdir build && cd build && \ - uv pip install 'cmake<4' setuptools numpy && \ + uv pip install 'cmake<4' 'setuptools<70' numpy && \ export PREFIX=/usr/local && CMAKE_ARGS="${CMAKE_ARGS} -DLLVM_ENABLE_PROJECTS=lld;libunwind;compiler-rt" \ CFLAGS="$(echo $CFLAGS | sed 's/-fno-plt //g')" \ CXXFLAGS="$(echo $CXXFLAGS | sed 's/-fno-plt //g')" \ @@ -193,27 +192,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \ cd opencv-python && \ python -m build --wheel --installer=uv --outdir /tmp/opencv-python/dist -# Build Outlines Core -FROM python-install AS outlines-core-builder +## Todo(r3hankhan123): Remove guidance-builder stage once vLLM upgrades to new version of llguidance that fixes s390x issues. See https://github.com/guidance-ai/llguidance/issues/330 +FROM python-install AS guidance-builder WORKDIR /tmp ENV CARGO_HOME=/root/.cargo ENV RUSTUP_HOME=/root/.rustup ENV PATH="$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH" -COPY requirements/common.txt /tmp/requirements/common.txt -ARG OUTLINES_CORE_VERSION RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \ --mount=type=bind,from=rust,source=/root/.rustup,target=/root/.rustup,rw \ - OUTLINES_CORE_VERSION=${OUTLINES_CORE_VERSION:-$(grep -E '^outlines_core\s*==\s*[0-9.]+' /tmp/requirements/common.txt | grep -Eo '[0-9.]+')} && \ - if [ -z "${OUTLINES_CORE_VERSION}" ]; then echo "ERROR: Could not determine outlines_core version"; exit 1; fi && \ - git clone https://github.com/dottxt-ai/outlines-core.git && \ - cd outlines-core && \ - git checkout tags/${OUTLINES_CORE_VERSION} && \ - sed -i "s/version = \"0.0.0\"/version = \"${OUTLINES_CORE_VERSION}\"/" Cargo.toml && \ + git clone https://github.com/guidance-ai/llguidance.git && \ + cd llguidance && \ + git checkout s390x-fix-v2 && \ uv pip install maturin && \ - python -m maturin build --release --out dist + python -m maturin build --release --out dist --compatibility linux -# Final build stage +# # Final build stage FROM python-install AS vllm-cpu ARG PYTHON_VERSION ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu" @@ -229,10 +223,12 @@ ENV PKG_CONFIG_PATH="/opt/rh/gcc-toolset-14/root/usr/lib64/pkgconfig:/usr/local/ ENV PATH="${VIRTUAL_ENV:+${VIRTUAL_ENV}/bin}:/opt/rh/gcc-toolset-14/root/usr/bin:/usr/local/bin:$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH" ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} +# Force pure Python protobuf to avoid s390x C++ extension crashes +ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python COPY . /workspace/vllm WORKDIR /workspace/vllm -RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.16,target=/numactl \ +RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.19,target=/numactl \ make -C /numactl install # Install dependencies, including PyTorch and Apache Arrow @@ -245,22 +241,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,from=numba-builder,source=/tmp/llvmlite/dist,target=/tmp/llvmlite-wheels/ \ --mount=type=bind,from=numba-builder,source=/tmp/numba/dist,target=/tmp/numba-wheels/ \ --mount=type=bind,from=opencv-builder,source=/tmp/opencv-python/dist,target=/tmp/opencv-wheels/ \ - --mount=type=bind,from=outlines-core-builder,source=/tmp/outlines-core/dist,target=/tmp/outlines-core/dist/ \ - ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/pyarrow-*.whl) && \ + --mount=type=bind,from=guidance-builder,source=/tmp/llguidance/dist,target=/tmp/guidance-wheels/ \ + ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/*.whl) && \ VISION_WHL_FILE=$(ls /tmp/vision-wheels/*.whl) && \ HF_XET_WHL_FILE=$(ls /tmp/hf-xet-wheels/*.whl) && \ LLVM_WHL_FILE=$(ls /tmp/llvmlite-wheels/*.whl) && \ NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \ OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \ - OUTLINES_CORE_WHL_FILE=$(ls /tmp/outlines-core/dist/*.whl) && \ - uv pip install -v \ - $ARROW_WHL_FILE \ + GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \ + uv pip install -v \ + $ARROW_WHL_FILE \ $VISION_WHL_FILE \ $HF_XET_WHL_FILE \ $LLVM_WHL_FILE \ $NUMBA_WHL_FILE \ $OPENCV_WHL_FILE \ - $OUTLINES_CORE_WHL_FILE \ + $GUIDANCE_WHL_FILE \ --index-strategy unsafe-best-match \ -r requirements/build/cpu.txt \ -r requirements/cpu.txt @@ -271,6 +267,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \ VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \ uv pip install "$(echo dist/*.whl)[tensorizer]" +# Remove protobuf C++ extension that crashes on s390x +RUN rm -rf /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/_upb/*.so \ + /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/protobuf/pyext/*.so 2>/dev/null || true + # setup non-root user for vllm RUN umask 002 && \ /usr/sbin/useradd --uid 2000 --gid 0 vllm && \ diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md index 3078c1537c0..1e36b431764 100644 --- a/docs/getting_started/installation/cpu.s390x.inc.md +++ b/docs/getting_started/installation/cpu.s390x.inc.md @@ -3,15 +3,15 @@ vLLM has experimental support for s390x architecture on IBM Z platform. For now, users must build from source to natively run on IBM Z platform. -Currently, the CPU implementation for s390x architecture supports FP32 datatype only. +Currently, the CPU implementation for s390x architecture supports FP32, BF16 and FP16. --8<-- [end:installation] --8<-- [start:requirements] - OS: `Linux` -- SDK: `gcc/g++ >= 12.3.0` or later with Command Line Tools +- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools - Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above. -- Build install python packages: `pyarrow`, `torch` and `torchvision` +- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless` --8<-- [end:requirements] --8<-- [start:set-up-using-python] @@ -24,13 +24,14 @@ Currently, there are no pre-built IBM Z CPU wheels. --8<-- [end:pre-built-wheels] --8<-- [start:build-wheel-from-source] -Install the following packages from the package manager before building the vLLM. For example on RHEL 9.4: +Install the following packages from the package manager before building the vLLM. For example on RHEL 9.6: ```bash dnf install -y \ - which procps findutils tar vim git gcc g++ make patch make cython zlib-devel \ + which procps findutils tar vim git gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \ libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \ - openssl-devel openblas openblas-devel wget autoconf automake libtool cmake numactl-devel + openssl-devel openblas openblas-devel autoconf automake libtool cmake numpy libsndfile \ + clang llvm-devel llvm-static clang-devel ``` Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation. @@ -43,13 +44,13 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \ Execute the following commands to build and install vLLM from source. !!! tip - Please build the following dependencies, `torchvision`, `pyarrow` from source before building vLLM. + Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM. ```bash - sed -i '/^torch/d' requirements/build/cuda.txt # remove torch from requirements/build/cuda.txt since we use nightly builds uv pip install -v \ + --extra-index-url https://download.pytorch.org/whl/cpu \ --torch-backend auto \ - -r requirements/build/cuda.txt \ + -r requirements/build/cpu.txt \ -r requirements/cpu.txt \ VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \ uv pip install dist/*.whl @@ -57,10 +58,9 @@ Execute the following commands to build and install vLLM from source. ??? console "pip" ```bash - sed -i '/^torch/d' requirements/build/cuda.txt # remove torch from requirements/build/cuda.txt since we use nightly builds pip install -v \ - --extra-index-url https://download.pytorch.org/whl/nightly/cpu \ - -r requirements/build/cuda.txt \ + --extra-index-url https://download.pytorch.org/whl/cpu \ + -r requirements/build/cpu.txt \ -r requirements/cpu.txt \ VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \ pip install dist/*.whl diff --git a/requirements/common.txt b/requirements/common.txt index 299ec734ff3..6e7fd90d023 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -19,7 +19,7 @@ pillow # Required for image processing prometheus-fastapi-instrumentator >= 7.0.0 tiktoken >= 0.6.0 # Required for DBRX tokenizer lm-format-enforcer == 0.11.3 -llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "s390x" or platform_machine == "ppc64le" +llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le" outlines_core == 0.2.11 # required for outlines backend disk cache diskcache == 5.6.3 From 7845379230c48c9103ab71682ee05977dfbbdfce Mon Sep 17 00:00:00 2001 From: realliujiaxu Date: Thu, 16 Apr 2026 13:48:00 +0800 Subject: [PATCH 044/150] [Bugfix] add support for 'num_attention_groups' in ModelArchConfigConvertorBase for Step3p5 (#39796) Signed-off-by: realliujiaxu --- tests/config/base_model_arch_groundtruth.json | 17 +++++++++++++++++ tests/config/test_model_arch_config.py | 1 + .../model_arch_config_convertor.py | 2 ++ 3 files changed, 20 insertions(+) diff --git a/tests/config/base_model_arch_groundtruth.json b/tests/config/base_model_arch_groundtruth.json index 81534886dcb..14e3f4f49d4 100644 --- a/tests/config/base_model_arch_groundtruth.json +++ b/tests/config/base_model_arch_groundtruth.json @@ -356,6 +356,23 @@ "is_multimodal_model": false, "dtype": "torch.float32" }, + "stepfun-ai/Step-3.5-Flash": { + "architectures": [ + "Step3p5ForCausalLM" + ], + "model_type": "step3p5", + "text_model_type": "step3p5", + "hidden_size": 4096, + "total_num_hidden_layers": 45, + "total_num_attention_heads": 64, + "head_size": 128, + "vocab_size": 128896, + "total_num_kv_heads": 8, + "num_experts": 288, + "is_deepseek_mla": false, + "is_multimodal_model": false, + "dtype": "torch.bfloat16" + }, "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16": { "architectures": [ "NemotronHForCausalLM" diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index fbae31331be..e172983b54f 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -16,6 +16,7 @@ BASE_TRUST_REMOTE_CODE_MODELS = { "nvidia/Llama-3_3-Nemotron-Super-49B-v1", "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", "XiaomiMiMo/MiMo-7B-RL", + "stepfun-ai/Step-3.5-Flash", # Excluded: Not available online right now # "FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1", "meituan-longcat/LongCat-Flash-Chat", diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index ea7096ae8bd..eb7d3eeda30 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -77,6 +77,8 @@ class ModelArchConfigConvertorBase: "num_key_value_heads", # For ChatGLM: "multi_query_group_num", + # For Step3p5: + "num_attention_groups", ] # For non-grouped-query attention models, the number of KV heads is # equal to the number of attention heads. From 2cdf86044d7e3bdcfbbb39a21a37b57ee9d4e65b Mon Sep 17 00:00:00 2001 From: Abhijit Roy <50175934+Roy214@users.noreply.github.com> Date: Thu, 16 Apr 2026 12:07:10 +0530 Subject: [PATCH 045/150] Add Jina Embeddings v5 model support (fixes #38633) (#39575) Signed-off-by: Abhijit Signed-off-by: wang.yuqi Co-authored-by: wang.yuqi --- docs/models/pooling_models/embed.md | 7 + tests/conftest.py | 6 + .../pooling_mteb_test/mteb_embed_utils.py | 30 +++- .../language/pooling_mteb_test/test_jina.py | 30 +++- tests/models/registry.py | 4 + vllm/model_executor/models/jina.py | 150 +++++++++++++++++- vllm/model_executor/models/registry.py | 1 + 7 files changed, 218 insertions(+), 10 deletions(-) diff --git a/docs/models/pooling_models/embed.md b/docs/models/pooling_models/embed.md index 8b3632a9f33..2f5d1a3fbe0 100644 --- a/docs/models/pooling_models/embed.md +++ b/docs/models/pooling_models/embed.md @@ -45,6 +45,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | | `GteModel` | Arctic-Embed-2.0-M | `Snowflake/snowflake-arctic-embed-m-v2.0`. | | | | `GteNewModel` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | | +| `JinaEmbeddingsV5Model`C | Qwen3-based with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small` (see note) | ✅︎ | ✅︎ | | `LlamaBidirectionalModel`C | Llama-based with bidirectional attention | `nvidia/llama-nemotron-embed-1b-v2`, etc. | ✅︎ | ✅︎ | | `LlamaModel`C, `LlamaForCausalLM`C, `MistralModel`C, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ | | `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | | @@ -73,6 +74,12 @@ You can compute pairwise similarity scores to build a similarity matrix using th !!! note `jinaai/jina-embeddings-v3` supports multiple tasks through LoRA, while vllm temporarily only supports text-matching tasks by merging LoRA weights. +!!! note + `jinaai/jina-embeddings-v5-text-small` ships with four task-specific LoRA adapters + (`retrieval`, `text-matching`, `classification`, `clustering`). vLLM merges the + selected adapter into the base weights at load time. Choose the task with + `--hf-overrides '{"jina_task": ""}'`; the default is `retrieval`. + ### Multimodal Models !!! note diff --git a/tests/conftest.py b/tests/conftest.py index bc657ff1ca7..4dbf3c8da15 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -364,6 +364,7 @@ class HfRunner: model_name: str, dtype: str = "auto", *, + revision: str | None = None, model_kwargs: dict[str, Any] | None = None, trust_remote_code: bool = True, is_sentence_transformer: bool = False, @@ -383,6 +384,7 @@ class HfRunner: self._init( model_name=model_name, dtype=dtype, + revision=revision, model_kwargs=model_kwargs, trust_remote_code=trust_remote_code, is_sentence_transformer=is_sentence_transformer, @@ -396,6 +398,7 @@ class HfRunner: model_name: str, dtype: str = "auto", *, + revision: str | None = None, model_kwargs: dict[str, Any] | None = None, trust_remote_code: bool = True, is_sentence_transformer: bool = False, @@ -437,6 +440,7 @@ class HfRunner: self.model = SentenceTransformer( model_name, + revision=revision, device=self.device, model_kwargs=model_kwargs, trust_remote_code=trust_remote_code, @@ -447,6 +451,7 @@ class HfRunner: self.model = CrossEncoder( model_name, + revision=revision, device=self.device, automodel_args=model_kwargs, trust_remote_code=trust_remote_code, @@ -456,6 +461,7 @@ class HfRunner: nn.Module, auto_cls.from_pretrained( model_name, + revision=revision, trust_remote_code=trust_remote_code, **model_kwargs, ), diff --git a/tests/models/language/pooling_mteb_test/mteb_embed_utils.py b/tests/models/language/pooling_mteb_test/mteb_embed_utils.py index 34b758d22ac..fc575c399d0 100644 --- a/tests/models/language/pooling_mteb_test/mteb_embed_utils.py +++ b/tests/models/language/pooling_mteb_test/mteb_embed_utils.py @@ -74,10 +74,25 @@ class MtebEmbedMixin(mteb.EncoderProtocol): return sim +class HfMtebEncoder(MtebEmbedMixin): + def __init__(self, model): + self.model = model + + def encode( + self, + inputs: DataLoader[mteb.types.BatchedInput], + *args, + **kwargs, + ) -> np.ndarray: + sentences = [text for batch in inputs for text in batch["text"]] + return self.model.encode(sentences) + + class VllmMtebEncoder(MtebEmbedMixin): - def __init__(self, vllm_model): + def __init__(self, vllm_model, prompt_prefix: str | None = None): self.llm = vllm_model self.rng = np.random.default_rng(seed=42) + self.prompt_prefix = prompt_prefix def encode( self, @@ -87,7 +102,11 @@ class VllmMtebEncoder(MtebEmbedMixin): ) -> np.ndarray: # Hoping to discover potential scheduling # issues by randomizing the order. - sentences = [text for batch in inputs for text in batch["text"]] + sentences = [ + self.prompt_prefix + text if self.prompt_prefix else text + for batch in inputs + for text in batch["text"] + ] r = self.rng.permutation(len(sentences)) sentences = [sentences[i] for i in r] outputs = self.llm.embed(sentences, use_tqdm=False) @@ -143,6 +162,7 @@ def mteb_test_embed_models( vllm_extra_kwargs=None, hf_model_callback=None, atol=MTEB_EMBED_TOL, + prompt_prefix: str | None = None, ): vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs) @@ -182,7 +202,7 @@ def mteb_test_embed_models( ) vllm_main_score = run_mteb_embed_task( - VllmMtebEncoder(vllm_model), MTEB_EMBED_TASKS + VllmMtebEncoder(vllm_model, prompt_prefix=prompt_prefix), MTEB_EMBED_TASKS ) vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype head_dtype = model_config.head_dtype @@ -210,7 +230,9 @@ def mteb_test_embed_models( if hf_model_callback is not None: hf_model_callback(hf_model) - st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS) + st_main_score = run_mteb_embed_task( + HfMtebEncoder(hf_model), MTEB_EMBED_TASKS + ) st_dtype = next(hf_model.model.parameters()).dtype # Check embeddings close to hf outputs diff --git a/tests/models/language/pooling_mteb_test/test_jina.py b/tests/models/language/pooling_mteb_test/test_jina.py index d75ec2a2ace..24aa3188f8b 100644 --- a/tests/models/language/pooling_mteb_test/test_jina.py +++ b/tests/models/language/pooling_mteb_test/test_jina.py @@ -28,7 +28,16 @@ EMBEDDING_MODELS = [ attn_type="encoder_only", is_prefix_caching_supported=False, is_chunked_prefill_supported=False, - ) + ), + EmbedModelInfo( + "jinaai/jina-embeddings-v5-text-small", + mteb_score=0.794535707854956, + architecture="JinaEmbeddingsV5Model", + seq_pooling_type="LAST", + attn_type="decoder", + is_prefix_caching_supported=True, + is_chunked_prefill_supported=True, + ), ] RERANK_MODELS = [ @@ -46,11 +55,18 @@ RERANK_MODELS = [ @pytest.mark.parametrize("model_info", EMBEDDING_MODELS) def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: + task = "retrieval" if "v5" in model_info.name else "text-matching" + prompt_prefix: str | None = "Document: " if "v5" in model_info.name else None + def hf_model_callback(model): - model.encode = partial(model.encode, task="text-matching") + model.encode = partial(model.encode, task=task) mteb_test_embed_models( - hf_runner, vllm_runner, model_info, hf_model_callback=hf_model_callback + hf_runner, + vllm_runner, + model_info, + hf_model_callback=hf_model_callback, + prompt_prefix=prompt_prefix, ) @@ -58,8 +74,10 @@ def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) - def test_embed_models_correctness( hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts ) -> None: + task = "retrieval" if "v5" in model_info.name else "text-matching" + def hf_model_callback(model): - model.encode = partial(model.encode, task="text-matching") + model.encode = partial(model.encode, task=task) correctness_test_embed_models( hf_runner, @@ -97,12 +115,14 @@ def test_matryoshka( # ST will strip the input texts, see test_embedding.py example_prompts = [str(s).strip() for s in example_prompts] + task = "retrieval" if "v5" in model_info.name else "text-matching" + with hf_runner( model_info.name, dtype=dtype, is_sentence_transformer=True, ) as hf_model: - hf_outputs = hf_model.encode(example_prompts, task="text-matching") + hf_outputs = hf_model.encode(example_prompts, task=task) hf_outputs = matryoshka_fy(hf_outputs, dimensions) with vllm_runner( diff --git a/tests/models/registry.py b/tests/models/registry.py index 956565e551b..f5968438cbd 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -609,6 +609,10 @@ _EMBEDDING_EXAMPLE_MODELS = { trust_remote_code=True, hf_overrides={"architectures": ["GteNewModel"]}, ), + "JinaEmbeddingsV5Model": _HfExamplesInfo( + "jinaai/jina-embeddings-v5-text-small", + trust_remote_code=True, + ), "LlamaModel": _HfExamplesInfo("llama", is_available_online=False), "LlamaBidirectionalModel": _HfExamplesInfo( "nvidia/llama-nemotron-embed-1b-v2", trust_remote_code=True diff --git a/vllm/model_executor/models/jina.py b/vllm/model_executor/models/jina.py index 980502191dd..2b07937df08 100644 --- a/vllm/model_executor/models/jina.py +++ b/vllm/model_executor/models/jina.py @@ -1,14 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from https://huggingface.co/jinaai/jina-reranker-v3/blob/main/modeling.py +import json +import logging +from collections import defaultdict from collections.abc import Iterable import torch +from safetensors.torch import load as safetensors_load from torch import nn from vllm.config import VllmConfig from vllm.sequence import IntermediateTensors from vllm.tasks import PoolingTask +from vllm.transformers_utils.repo_utils import get_hf_file_bytes from vllm.v1.pool.metadata import PoolingMetadata from ..layers.pooler import DispatchPooler @@ -18,9 +23,12 @@ from ..layers.pooler.tokwise import ( TokenPoolingMethodOutputItem, ) from .interfaces import SupportsLateInteraction -from .qwen3 import Qwen3Model +from .interfaces_base import VllmModelForPooling +from .qwen3 import Qwen3ForCausalLM, Qwen3Model from .utils import AutoWeightsLoader, maybe_prefix +logger = logging.getLogger(__name__) + class JinaForRanking(nn.Module, SupportsLateInteraction): is_pooling_model = True @@ -108,3 +116,143 @@ class JinaForRankingPool(StepPool): embeds_list.append(embeds) return embeds_list + + +# jina-embeddings-v5-text-small wraps Qwen3-0.6B-Base with four task-specific +# LoRA adapters. This implementation merges the selected adapter into the base +# weights at load time to avoid any runtime dependency on peft. +# +# Task selection: +# Pass --hf-overrides '{"jina_task": "retrieval"}' to select one of: +# retrieval (default), text-matching, classification, clustering. + +_DEFAULT_TASK = "retrieval" +_SUPPORTED_TASKS = {"retrieval", "text-matching", "classification", "clustering"} + + +def _load_adapter( + model: str, + task: str, + revision: str | None, +) -> tuple[dict, dict[str, torch.Tensor]] | None: + """Load adapter config and weights from a local path or HF repo. + + Returns (adapter_config, adapter_weights) or None if not found. + """ + config_bytes = get_hf_file_bytes( + f"adapters/{task}/adapter_config.json", + model, + revision, + ) + if config_bytes is None: + return None + + adapter_config = json.loads(config_bytes) + + weights_bytes = get_hf_file_bytes( + f"adapters/{task}/adapter_model.safetensors", + model, + revision, + ) + if weights_bytes is None: + return None + + adapter_weights = safetensors_load(weights_bytes) + return adapter_config, adapter_weights + + +def _build_lora_pairs(adapter_weights: dict) -> dict: + """Group raw adapter tensors into {base_key: {"A": tensor, "B": tensor}} pairs. + + Transforms adapter keys like: + base_model.model.layers.0.self_attn.q_proj.lora_A.weight + Into base keys like: + layers.0.self_attn.q_proj.weight + """ + lora_pairs = defaultdict(dict) + for key, tensor in adapter_weights.items(): + clean_key = key + if clean_key.startswith("base_model.model."): + clean_key = clean_key[len("base_model.model.") :] + + if ".lora_A." in clean_key: + base_key = clean_key.split(".lora_A.")[0] + ".weight" + lora_pairs[base_key]["A"] = tensor + elif ".lora_B." in clean_key: + base_key = clean_key.split(".lora_B.")[0] + ".weight" + lora_pairs[base_key]["B"] = tensor + + return dict(lora_pairs) + + +class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling): + """Jina Embeddings V5 with task-specific LoRA adapters merged at load time. + + Extends Qwen3ForCausalLM (the underlying architecture) and declares itself + as a pooling model so that as_embedding_model() does not wrap it. + """ + + is_pooling_model = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + + self._model_name = vllm_config.model_config.model + self._revision = vllm_config.model_config.revision + + self._task = getattr( + vllm_config.model_config.hf_config, "jina_task", _DEFAULT_TASK + ) + if self._task not in _SUPPORTED_TASKS: + logger.warning( + "Unknown jina_task=%r. Falling back to %r.", + self._task, + _DEFAULT_TASK, + ) + self._task = _DEFAULT_TASK + + pooler_config = vllm_config.model_config.pooler_config + assert pooler_config is not None + self.pooler = DispatchPooler.for_embedding(pooler_config) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + lora_pairs: dict = {} + scaling = 1.0 + + result = _load_adapter(self._model_name, self._task, self._revision) + if result is None: + logger.warning( + "No adapter found for task %r in %r. Loading raw base weights.", + self._task, + self._model_name, + ) + else: + adapter_config, adapter_weights = result + scaling = adapter_config["lora_alpha"] / adapter_config["r"] + lora_pairs = _build_lora_pairs(adapter_weights) + logger.info( + "Loaded %d adapter tensors for task %r (scaling=%.4f, %d LoRA pairs)", + len(adapter_weights), + self._task, + scaling, + len(lora_pairs), + ) + + def _merge_weights( + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Iterable[tuple[str, torch.Tensor]]: + for name, tensor in weights: + clean_name = name + if clean_name.startswith("model."): + clean_name = clean_name[len("model.") :] + + if clean_name in lora_pairs: + pair = lora_pairs[clean_name] + if "A" in pair and "B" in pair: + lora_A = pair["A"].to(device=tensor.device, dtype=tensor.dtype) + lora_B = pair["B"].to(device=tensor.device, dtype=tensor.dtype) + tensor = tensor + (lora_B @ lora_A) * scaling + yield name, tensor + + loaded = self.model.load_weights(_merge_weights(weights)) + return {f"model.{name}" for name in loaded} diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 01a767b9f32..b3cebd5100b 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -227,6 +227,7 @@ _EMBEDDING_MODELS = { "GritLM": ("gritlm", "GritLM"), "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"), "GteNewModel": ("bert_with_rope", "GteNewModel"), + "JinaEmbeddingsV5Model": ("jina", "JinaEmbeddingsV5Model"), "LlamaBidirectionalModel": ("llama", "LlamaBidirectionalModel"), "LlamaModel": ("llama", "LlamaForCausalLM"), **{ From f4ddaf8cf7f7daf1233acbfd4db97ca107a0a5b7 Mon Sep 17 00:00:00 2001 From: Xinyu Chen Date: Thu, 16 Apr 2026 14:42:07 +0800 Subject: [PATCH 046/150] [XPU] use spawn multiproc method on xpu (#39671) Signed-off-by: Xinyu Chen Co-authored-by: Kunshang Ji --- vllm/platforms/xpu.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index f9a4f8bd732..aa673419730 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -228,6 +228,10 @@ class XPUPlatform(Platform): # ref. https://openucx.readthedocs.io/en/master/faq.html os.environ["UCX_MEMTYPE_CACHE"] = "n" + # spawn is the only supported multiprocessing method on XPU + if "VLLM_WORKER_MULTIPROC_METHOD" not in os.environ: + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + @classmethod def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: super().update_block_size_for_backend(vllm_config) From 8d7c9628337aebe0f78239de0b385e82c631abc6 Mon Sep 17 00:00:00 2001 From: Tim Messerschmidt Date: Thu, 16 Apr 2026 09:18:32 +0200 Subject: [PATCH 047/150] [Bugfix] Accept **kwargs in MiniMaxM2Parser.__init__() (#39861) Signed-off-by: Tim Messerschmidt --- vllm/parser/minimax_m2_parser.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/vllm/parser/minimax_m2_parser.py b/vllm/parser/minimax_m2_parser.py index c3211429464..34aaa726844 100644 --- a/vllm/parser/minimax_m2_parser.py +++ b/vllm/parser/minimax_m2_parser.py @@ -43,11 +43,17 @@ class MiniMaxM2Parser(DelegatingParser): reasoning_parser_cls = MiniMaxM2ReasoningParser tool_parser_cls = MinimaxM2ToolParser - def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): - super().__init__(tokenizer) + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + *args, + **kwargs, + ): + super().__init__(tokenizer, *args, **kwargs) # Initialize the underlying parsers - self._reasoning_parser = MiniMaxM2ReasoningParser(tokenizer) + self._reasoning_parser = MiniMaxM2ReasoningParser(tokenizer, *args, **kwargs) self._tool_parser = MinimaxM2ToolParser(tokenizer, tools) logger.debug( From 10e49d263854daf6cf63472b9cd2039196022a59 Mon Sep 17 00:00:00 2001 From: Simon Mo Date: Thu, 16 Apr 2026 00:22:03 -0700 Subject: [PATCH 048/150] [Docs] Update PR template to remove release notes google docs (#39982) Signed-off-by: Simon Mo --- .github/PULL_REQUEST_TEMPLATE.md | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 8043df65d55..8a3934670e4 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -15,7 +15,6 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT - [ ] The test plan, such as providing test command. - [ ] The test results, such as pasting the results comparison before and after, or e2e results - [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. -- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft in the [Google Doc](https://docs.google.com/document/d/1YyVqrgX4gHTtrstbq8oWUImOyPCKSGnJ7xtTpmXzlRs/edit?tab=t.0). **BEFORE SUBMITTING, PLEASE READ ** (anything written below this line will be removed by GitHub Actions) From 98700c6105b0313d924126cba428219111ee8f6f Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Thu, 16 Apr 2026 12:06:51 +0300 Subject: [PATCH 049/150] Fix #33773: Replace unconditional pandas import with PlaceholderModule (#39990) Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> --- vllm/_aiter_ops.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index d71e210fb51..55d6d1297a8 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -3,18 +3,23 @@ import functools from collections.abc import Callable -import pandas as pd import torch from torch._ops import OpOverload import vllm.envs as envs from vllm.platforms import current_platform +from vllm.utils.import_utils import PlaceholderModule from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( rocm_aiter_sparse_attn_indexer, rocm_aiter_sparse_attn_indexer_fake, ) +try: + import pandas as pd +except ImportError: + pd = PlaceholderModule("pandas") + # fp8_dtype is not cached. # on ROCm the fp8_dtype always calls is_fp8_fnuz # which is a host op, so we cache it once here. From 17d87168d27c2d2a99c544b57ea046629bcdc48f Mon Sep 17 00:00:00 2001 From: lalit10 Date: Thu, 16 Apr 2026 02:16:06 -0700 Subject: [PATCH 050/150] [Model] Use mm_features for Keye-VL and Keye-1.5-VL M-RoPE (#39869) Signed-off-by: Lalit Laxminarayan Bangad --- tests/model_executor/test_keye_mrope.py | 145 ++++++++++++++++ tests/model_executor/test_keye_vl1_5_mrope.py | 145 ++++++++++++++++ vllm/model_executor/models/keye.py | 159 +++++++++--------- vllm/model_executor/models/keye_vl1_5.py | 134 ++++++--------- 4 files changed, 425 insertions(+), 158 deletions(-) create mode 100644 tests/model_executor/test_keye_mrope.py create mode 100644 tests/model_executor/test_keye_vl1_5_mrope.py diff --git a/tests/model_executor/test_keye_mrope.py b/tests/model_executor/test_keye_mrope.py new file mode 100644 index 00000000000..2289a59e2bf --- /dev/null +++ b/tests/model_executor/test_keye_mrope.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass, field + +import pytest +import torch + +from vllm.model_executor.models.keye import KeyeForConditionalGeneration +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalFieldElem, + MultiModalKwargsItem, + PlaceholderRange, +) + +pytestmark = pytest.mark.skip_global_cleanup + + +@pytest.fixture(autouse=True, scope="module") +def _force_cpu_default_device(): + original = torch.get_default_device() + torch.set_default_device("cpu") + yield + torch.set_default_device(original) + + +@dataclass +class DummyVisionConfig: + spatial_merge_size: int = 2 + + +@dataclass +class DummyConfig: + vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig) + + +def make_model(config: DummyConfig) -> KeyeForConditionalGeneration: + model = object.__new__(KeyeForConditionalGeneration) + model.config = config + return model + + +def make_mm_feature( + *, + modality: str, + offset: int, + length: int, + grid_thw: tuple[int, int, int] | list[tuple[int, int, int]], +) -> MultiModalFeatureSpec: + field_name = "image_grid_thw" if modality == "image" else "video_grid_thw" + return MultiModalFeatureSpec( + data=MultiModalKwargsItem( + { + field_name: MultiModalFieldElem( + data=torch.tensor(grid_thw), + field=None, # HACK. + ), + } + ), + modality=modality, + identifier="DUMMY", + mm_position=PlaceholderRange(offset=offset, length=length), + ) + + +def test_get_mrope_input_positions_text_only(): + model = make_model(DummyConfig()) + + positions, delta = model.get_mrope_input_positions( + input_tokens=[11, 12, 13, 14, 15], + mm_features=[], + ) + + expected = torch.tensor( + [ + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + ] + ) + + assert torch.equal(positions, expected) + assert delta == 0 + + +def test_get_mrope_input_positions_single_image(): + model = make_model(DummyConfig()) + mm_features = [ + make_mm_feature( + modality="image", + offset=1, + length=4, + grid_thw=(1, 4, 4), + ) + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=[10, 20, 21, 22, 23, 30, 31], + mm_features=mm_features, + ) + + expected = torch.tensor( + [ + [0, 1, 1, 1, 1, 3, 4], + [0, 1, 1, 2, 2, 3, 4], + [0, 1, 2, 1, 2, 3, 4], + ] + ) + + assert torch.equal(positions, expected) + assert delta == -2 + + +def test_get_mrope_input_positions_interleaved_image_and_video(): + model = make_model(DummyConfig()) + mm_features = [ + make_mm_feature( + modality="image", + offset=1, + length=4, + grid_thw=(1, 4, 4), + ), + make_mm_feature( + modality="video", + offset=7, + length=4, + grid_thw=[(2, 4, 2)], + ), + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 42, 43, 50, 51], + mm_features=mm_features, + ) + + expected = torch.tensor( + [ + [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 7, 9, 10], + [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 7, 9, 10], + ] + ) + + assert torch.equal(positions, expected) + assert delta == -2 diff --git a/tests/model_executor/test_keye_vl1_5_mrope.py b/tests/model_executor/test_keye_vl1_5_mrope.py new file mode 100644 index 00000000000..de3488bc922 --- /dev/null +++ b/tests/model_executor/test_keye_vl1_5_mrope.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass, field + +import pytest +import torch + +from vllm.model_executor.models.keye_vl1_5 import KeyeVL1_5ForConditionalGeneration +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalFieldElem, + MultiModalKwargsItem, + PlaceholderRange, +) + +pytestmark = pytest.mark.skip_global_cleanup + + +@pytest.fixture(autouse=True, scope="module") +def _force_cpu_default_device(): + original = torch.get_default_device() + torch.set_default_device("cpu") + yield + torch.set_default_device(original) + + +@dataclass +class DummyVisionConfig: + spatial_merge_size: int = 2 + + +@dataclass +class DummyConfig: + vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig) + + +def make_model(config: DummyConfig) -> KeyeVL1_5ForConditionalGeneration: + model = object.__new__(KeyeVL1_5ForConditionalGeneration) + model.config = config + return model + + +def make_mm_feature( + *, + modality: str, + offset: int, + length: int, + grid_thw: tuple[int, int, int] | list[tuple[int, int, int]], + is_embed: list[bool] | None = None, +) -> MultiModalFeatureSpec: + field_name = "image_grid_thw" if modality == "image" else "video_grid_thw" + return MultiModalFeatureSpec( + data=MultiModalKwargsItem( + { + field_name: MultiModalFieldElem( + data=torch.tensor(grid_thw), + field=None, # HACK. + ), + } + ), + modality=modality, + identifier="DUMMY", + mm_position=PlaceholderRange( + offset=offset, + length=length, + is_embed=None if is_embed is None else torch.tensor(is_embed), + ), + ) + + +def test_get_mrope_input_positions_text_only(): + model = make_model(DummyConfig()) + + positions, delta = model.get_mrope_input_positions( + input_tokens=[11, 12, 13, 14, 15], + mm_features=[], + ) + + expected = torch.tensor( + [ + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + ] + ) + + assert torch.equal(positions, expected) + assert delta == 0 + + +def test_get_mrope_input_positions_single_image(): + model = make_model(DummyConfig()) + mm_features = [ + make_mm_feature( + modality="image", + offset=1, + length=4, + grid_thw=(1, 4, 4), + ) + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=[10, 20, 21, 22, 23, 30, 31], + mm_features=mm_features, + ) + + expected = torch.tensor( + [ + [0, 1, 1, 1, 1, 3, 4], + [0, 1, 1, 2, 2, 3, 4], + [0, 1, 2, 1, 2, 3, 4], + ] + ) + + assert torch.equal(positions, expected) + assert delta == -2 + + +def test_get_mrope_input_positions_video_uses_embed_ranges(): + model = make_model(DummyConfig()) + mm_features = [ + make_mm_feature( + modality="video", + offset=1, + length=8, + grid_thw=[(2, 4, 2)], + is_embed=[False, False, True, True, False, False, True, True], + ) + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=[10, 101, 102, 20, 21, 103, 104, 30, 31, 40, 41], + mm_features=mm_features, + ) + + expected = torch.tensor( + [ + [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10], + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10], + ] + ) + + assert torch.equal(positions, expected) + assert delta == 0 diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index a987c89ae09..7166f725774 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from abc import abstractmethod -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from functools import partial from typing import Annotated, Any, Literal, TypeAlias, TypeVar @@ -1595,91 +1595,92 @@ class KeyeForConditionalGeneration( self._process_video_embeds(video_type, video_grid_thw, pixel_values_videos) ) + @staticmethod + def _split_video_grid_thw( + grid_thw: torch.Tensor | list[list[int]] | list[int], + ) -> list[list[int]]: + """ + Split video grid_thw along the t dimension into per-frame rows. + + This preserves Keye's current M-RoPE behavior, where a video is emitted + as consecutive frame-level multimodal blocks rather than a single block + spanning the whole video. + """ + if isinstance(grid_thw, list): + if len(grid_thw) == 0: + return [] + if isinstance(grid_thw[0], int): + grid_thw = torch.tensor([grid_thw], dtype=torch.long) + else: + grid_thw = torch.tensor(grid_thw, dtype=torch.long) + elif grid_thw.ndim == 1: + grid_thw = grid_thw.unsqueeze(0) + + if grid_thw.numel() == 0: + return [] + + t, hw = grid_thw[:, 0], grid_thw[:, 1:] + ones = torch.ones_like(hw[:, :1]) + out = torch.cat([ones, hw], dim=1).repeat_interleave(t, dim=0) + return out.tolist() + + def iter_mm_grid_thw( + self, mm_features: list[MultiModalFeatureSpec] + ) -> Iterator[tuple[int, int, int, int]]: + spatial_merge_size = self.config.vision_config.spatial_merge_size + + for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): + if mm_feature.data is None: + raise ValueError("M-RoPE calculation requires multimodal feature data") + + if mm_feature.modality == "image": + grid_thw = mm_feature.data["image_grid_thw"].data + if isinstance(grid_thw, torch.Tensor): + if grid_thw.ndim == 2: + assert grid_thw.shape[0] == 1 + t, h, w = grid_thw[0].tolist() + else: + t, h, w = grid_thw.tolist() + else: + if isinstance(grid_thw[0], list): + assert len(grid_thw) == 1 + t, h, w = grid_thw[0] + else: + t, h, w = grid_thw + + yield ( + mm_feature.mm_position.offset, + t, + h // spatial_merge_size, + w // spatial_merge_size, + ) + elif mm_feature.modality == "video": + current_offset = mm_feature.mm_position.offset + for t, h, w in self._split_video_grid_thw( + mm_feature.data["video_grid_thw"].data + ): + llm_grid_h = h // spatial_merge_size + llm_grid_w = w // spatial_merge_size + yield (current_offset, t, llm_grid_h, llm_grid_w) + current_offset += t * llm_grid_h * llm_grid_w + else: + raise ValueError(f"Unsupported modality: {mm_feature.modality}") + def get_mrope_input_positions( self, input_tokens: list[int], mm_features: list[MultiModalFeatureSpec], ) -> tuple[torch.Tensor, int]: - kwargs = MultiModalFeatureSpec.gather_kwargs( - mm_features, - {"image_grid_thw", "video_grid_thw"}, - ) - image_grid_thw = [item.tolist() for item in kwargs.get("image_grid_thw", [])] - video_grid_thw = [item.tolist() for item in kwargs.get("video_grid_thw", [])] - - if isinstance(video_grid_thw, list) and len(video_grid_thw) > 0: - video_grid_thw = video_grid_thw[0] - - def split_thw(grid_thw: torch.Tensor | list[int]) -> list[list[int]]: - """ - Split grid_thw along the t dimension. - - Args: - grid_thw: shape [N, 3] tensor or nested list of [t, h, w]. - - Returns: - List of [1, h, w] rows, repeated t times for each original row. - """ - - if isinstance(grid_thw, list): - grid_thw = torch.tensor(grid_thw, dtype=torch.long) - - if grid_thw.numel() == 0: - return [] - - t, hw = grid_thw[:, 0], grid_thw[:, 1:] - ones = torch.ones_like(hw[:, :1]) # [N,1] - out = torch.cat([ones, hw], dim=1).repeat_interleave(t, dim=0) - return out.tolist() - - video_grid_thw = split_thw(video_grid_thw) - - hf_config = self.config - image_token_id = hf_config.image_token_id - video_token_id = hf_config.video_token_id - spatial_merge_size = hf_config.vision_config.spatial_merge_size - - image_nums = len(image_grid_thw) - frame_nums = len(video_grid_thw) llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_frames = image_nums, frame_nums - image_index, video_index = 0, 0 - for _ in range(image_nums + frame_nums): - if remain_images > 0: - try: - ed_image = input_tokens.index(image_token_id, st) - except ValueError: - ed_image = len(input_tokens) + 1 - else: - ed_image = len(input_tokens) + 1 - if remain_frames > 0: - try: - ed_video = input_tokens.index(video_token_id, st) - except ValueError: - ed_video = len(input_tokens) + 1 - else: - ed_video = len(input_tokens) + 1 - - if ed_image < ed_video: - t, h, w = image_grid_thw[image_index] - image_index += 1 - remain_images -= 1 - ed = ed_image - else: - t, h, w = video_grid_thw[video_index] - video_index += 1 - remain_frames -= 1 - ed = ed_video - - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - text_len = ed - st + for ( + offset, + llm_grid_t, + llm_grid_h, + llm_grid_w, + ) in self.iter_mm_grid_thw(mm_features): + text_len = offset - st st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 llm_pos_ids_list.append( @@ -1711,7 +1712,7 @@ class KeyeForConditionalGeneration( llm_pos_ids_list.append( torch.stack([t_index, h_index, w_index]) + text_len + st_idx ) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w + st = offset + llm_grid_t * llm_grid_h * llm_grid_w if st < len(input_tokens): st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 diff --git a/vllm/model_executor/models/keye_vl1_5.py b/vllm/model_executor/models/keye_vl1_5.py index bc33f5d7d72..a8400b8d17e 100644 --- a/vllm/model_executor/models/keye_vl1_5.py +++ b/vllm/model_executor/models/keye_vl1_5.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import itertools -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from functools import partial from typing import Annotated, Any, Literal, TypeAlias @@ -608,91 +608,67 @@ class KeyeVL1_5ForConditionalGeneration( new_video_embeds.append(video_embeds[start:end]) return tuple(new_video_embeds) + def iter_mm_grid_thw( + self, mm_features: list[MultiModalFeatureSpec] + ) -> Iterator[tuple[int, int, int, int]]: + spatial_merge_size = self.config.vision_config.spatial_merge_size + + for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): + if mm_feature.data is None: + raise ValueError("M-RoPE calculation requires multimodal feature data") + + embed_ranges = mm_feature.mm_position.extract_embeds_range() + if mm_feature.modality == "image": + assert len(embed_ranges) == 1 + grid_thw = mm_feature.data["image_grid_thw"].data + if isinstance(grid_thw, torch.Tensor): + if grid_thw.ndim == 2: + assert grid_thw.shape[0] == 1 + t, h, w = grid_thw[0].tolist() + else: + t, h, w = grid_thw.tolist() + else: + if isinstance(grid_thw[0], list): + assert len(grid_thw) == 1 + t, h, w = grid_thw[0] + else: + t, h, w = grid_thw + + yield ( + embed_ranges[0][0], + t, + h // spatial_merge_size, + w // spatial_merge_size, + ) + elif mm_feature.modality == "video": + split_video_grids = split_thw(mm_feature.data["video_grid_thw"].data) + assert len(embed_ranges) == split_video_grids.shape[0] + for (start_idx, end_idx), (t, h, w) in zip( + embed_ranges, split_video_grids.tolist() + ): + llm_grid_h = h // spatial_merge_size + llm_grid_w = w // spatial_merge_size + num_mm_tokens = t * llm_grid_h * llm_grid_w + assert end_idx - start_idx + 1 == num_mm_tokens + yield (start_idx, t, llm_grid_h, llm_grid_w) + else: + raise ValueError(f"Unsupported modality: {mm_feature.modality}") + def get_mrope_input_positions( self, input_tokens: list[int], mm_features: list[MultiModalFeatureSpec], ) -> tuple[torch.Tensor, int]: - kwargs = MultiModalFeatureSpec.gather_kwargs( - mm_features, - {"image_grid_thw", "video_grid_thw"}, - ) - image_grid_thw = [item.tolist() for item in kwargs.get("image_grid_thw", [])] - video_grid_thw = [item.tolist() for item in kwargs.get("video_grid_thw", [])] - - if isinstance(video_grid_thw, list) and len(video_grid_thw) > 0: - video_grid_thw = video_grid_thw[0] - - def split_thw(grid_thw: torch.Tensor | list[int]) -> list[list[int]]: - """ - Split grid_thw along the t dimension. - - Args: - grid_thw: shape [N, 3] tensor or nested list of [t, h, w]. - - Returns: - List of [1, h, w] rows, repeated t times for each original row. - """ - - if isinstance(grid_thw, list): - grid_thw = torch.tensor(grid_thw, dtype=torch.long) - - if grid_thw.numel() == 0: - return [] - - t, hw = grid_thw[:, 0], grid_thw[:, 1:] - ones = torch.ones_like(hw[:, :1]) # [N,1] - out = torch.cat([ones, hw], dim=1).repeat_interleave(t, dim=0) - return out.tolist() - - video_grid_thw = split_thw(video_grid_thw) - - hf_config = self.config - image_token_id = hf_config.image_token_id - video_token_id = hf_config.video_token_id - spatial_merge_size = hf_config.vision_config.spatial_merge_size - - image_nums = len(image_grid_thw) - frame_nums = len(video_grid_thw) llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_frames = image_nums, frame_nums - image_index, video_index = 0, 0 - for _ in range(image_nums + frame_nums): - if remain_images > 0: - try: - ed_image = input_tokens.index(image_token_id, st) - except ValueError: - ed_image = len(input_tokens) + 1 - else: - ed_image = len(input_tokens) + 1 - if remain_frames > 0: - try: - ed_video = input_tokens.index(video_token_id, st) - except ValueError: - ed_video = len(input_tokens) + 1 - else: - ed_video = len(input_tokens) + 1 - - if ed_image < ed_video: - t, h, w = image_grid_thw[image_index] - image_index += 1 - remain_images -= 1 - ed = ed_image - else: - t, h, w = video_grid_thw[video_index] - video_index += 1 - remain_frames -= 1 - ed = ed_video - - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - text_len = ed - st + for ( + offset, + llm_grid_t, + llm_grid_h, + llm_grid_w, + ) in self.iter_mm_grid_thw(mm_features): + text_len = offset - st st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 llm_pos_ids_list.append( @@ -724,7 +700,7 @@ class KeyeVL1_5ForConditionalGeneration( llm_pos_ids_list.append( torch.stack([t_index, h_index, w_index]) + text_len + st_idx ) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w + st = offset + llm_grid_t * llm_grid_h * llm_grid_w if st < len(input_tokens): st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 From 9965f501a89204769a53c86cdee2528947373747 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 16 Apr 2026 12:53:21 +0200 Subject: [PATCH 051/150] [Nixl] Bump Nixl version to 0.10.1 (#39922) Signed-off-by: NickLucche --- requirements/kv_connectors.txt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index 700213feda1..ae2173a9e8d 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -1,5 +1,5 @@ lmcache >= 0.3.9 -nixl[cu13] >= 0.7.1, < 0.10.0 # Required for disaggregated prefill -nixl-cu12 >= 0.7.1, < 0.10.0 -nixl-cu13 >= 0.7.1, < 0.10.0 +nixl[cu13] >= 0.7.1, <= 0.10.1 # Required for disaggregated prefill +nixl-cu12 >= 0.7.1, <= 0.10.1 +nixl-cu13 >= 0.7.1, <= 0.10.1 mooncake-transfer-engine >= 0.3.8 From edc3648966e30e8bf5f34edeee50027ddd79dc16 Mon Sep 17 00:00:00 2001 From: Yanan Cao Date: Thu, 16 Apr 2026 04:41:26 -0700 Subject: [PATCH 052/150] [Kernel][Helion] Fix inductor fusion of Helion HOP (#39944) Signed-off-by: Yanan Cao Co-authored-by: Claude Opus 4.6 (1M context) --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/kernels.yaml | 2 +- setup.py | 2 +- tests/kernels/helion/test_register.py | 48 +++++++++++++++++++++++++++ vllm/kernels/helion/register.py | 11 ++++-- 5 files changed, 59 insertions(+), 6 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index dda5d4064c3..847d8762da2 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -769,7 +769,7 @@ steps: - tests/kernels/helion/ - vllm/platforms/rocm.py commands: - - pip install helion==0.3.3 + - pip install helion==1.0.0 - pytest -v -s kernels/helion/ diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 8b9765130ae..3e7376e4134 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -155,7 +155,7 @@ steps: - vllm/utils/import_utils.py - tests/kernels/helion/ commands: - - pip install helion==0.3.3 + - pip install helion==1.0.0 - pytest -v -s kernels/helion/ diff --git a/setup.py b/setup.py index b0cca73bb91..af52cb58dfc 100644 --- a/setup.py +++ b/setup.py @@ -1104,7 +1104,7 @@ setup( # NOTE: When updating helion version, also update CI files: # - .buildkite/test_areas/kernels.yaml # - .buildkite/test-amd.yaml - "helion": ["helion==0.3.3"], + "helion": ["helion==1.0.0"], # Optional deps for gRPC server (vllm serve --grpc) "grpc": ["smg-grpc-servicer[vllm] >= 0.5.0"], # Optional deps for OpenTelemetry tracing diff --git a/tests/kernels/helion/test_register.py b/tests/kernels/helion/test_register.py index c7f93993c57..bad3017c5c9 100644 --- a/tests/kernels/helion/test_register.py +++ b/tests/kernels/helion/test_register.py @@ -36,9 +36,11 @@ from vllm.kernels.helion.register import ( ) if _HOP_AVAILABLE: + from helion._compat import supports_torch_compile_fusion from helion._compiler._dynamo.higher_order_ops import ( helion_kernel_wrapper_mutation, ) + from torch._inductor.utils import run_and_get_code def _add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: @@ -1003,3 +1005,49 @@ class TestTorchCompileHOP: "Compiled execution result doesn't match eager execution. " f"Max difference: {torch.max(torch.abs(compiled_result - eager_result))}" ) + + @pytest.mark.skipif( + not (_HOP_AVAILABLE and supports_torch_compile_fusion()), + reason="Requires PyTorch with Helion inductor fusion support", + ) + def test_inductor_backend_compiles_helion_hop(self): + """Test torch.compile with inductor backend and Helion fusion enabled.""" + + configs = {"default": helion.Config(block_sizes=[4, 4])} + + with dummy_kernel_registry(configs=configs) as register: + add_helion_kernel = register( + op_name="test_inductor_add_kernel", + config_picker=lambda args, keys: "default", + helion_settings=helion.Settings( + torch_compile_fusion=True, static_shapes=False + ), + )(_add_kernel) + + def f(x, y): + x = x * 2.0 + y = y + 1.0 + out = add_helion_kernel(x, y) + return out.relu() + + torch._dynamo.reset() + compiled_f = torch.compile(f, backend="inductor", fullgraph=True) + + x = torch.randn(4, 4, device="cuda") + y = torch.randn(4, 4, device="cuda") + + compiled_result, source_codes = run_and_get_code(compiled_f, x, y) + eager_result = f(x, y) + + assert torch.allclose(compiled_result, eager_result, atol=1e-5, rtol=1e-5), ( + "Inductor-compiled result doesn't match eager execution. " + f"Max difference: {torch.max(torch.abs(compiled_result - eager_result))}" + ) + + # With fusion enabled, prologue/epilogue ops should be fused into + # a single triton kernel rather than generating separate kernels. + kernel_count = sum(code.count("@triton.jit") for code in source_codes) + assert kernel_count == 1, ( + f"Expected 1 fused triton kernel, got {kernel_count}. " + "Prologue/epilogue ops were not fused into the Helion kernel." + ) diff --git a/vllm/kernels/helion/register.py b/vllm/kernels/helion/register.py index 1557a36b275..080c4cdda0f 100644 --- a/vllm/kernels/helion/register.py +++ b/vllm/kernels/helion/register.py @@ -63,8 +63,10 @@ from helion.runtime.settings import default_autotuner_fn _HOP_AVAILABLE = requires_torch_version("2.11") if _HOP_AVAILABLE: + from helion._compat import supports_torch_compile_fusion from helion._compiler._dynamo.higher_order_ops import helion_kernel_side_table from helion._compiler._dynamo.variables import HelionKernelVariable + from helion.runtime.kernel import Kernel from torch._dynamo.guards import GuardBuilder from torch._dynamo.variables.builder import VariableBuilder @@ -475,19 +477,22 @@ if _HOP_AVAILABLE: """Register HelionKernelWrapper with Dynamo's VariableBuilder. When Dynamo encounters a HelionKernelWrapper during tracing, this - extracts the underlying Helion Kernel, registers it in the side table, - and returns Helion's own HelionKernelVariable to handle HOP emission. + extracts the underlying Helion Kernel and delegates to Helion's own + registered Kernel handler, which handles HOP emission, side table + registration, and inductor lowering setup. """ def wrap_helion_kernel_wrapper( builder: VariableBuilder, value: HelionKernelWrapper ): kernel = value.get_configured_op()._decorated_kernel + if supports_torch_compile_fusion(): + helion_handler = VariableBuilder._type_dispatch()[Kernel] + return helion_handler(builder, kernel) kernel_idx = helion_kernel_side_table.add_kernel(kernel) builder.install_guards(GuardBuilder.ID_MATCH) return HelionKernelVariable(kernel, kernel_idx, source=builder.source) - # Register with Dynamo's type dispatch system dispatch = VariableBuilder._type_dispatch() dispatch[HelionKernelWrapper] = wrap_helion_kernel_wrapper From 4269b794091c057491b2ee5ac855d5a268959cb7 Mon Sep 17 00:00:00 2001 From: grYe99 <35474918+grYe99@users.noreply.github.com> Date: Thu, 16 Apr 2026 21:14:00 +0800 Subject: [PATCH 053/150] [Model] Use mm_features to compute mrope positions for PaddleOCR-VL (#39888) Signed-off-by: grYe99 Co-authored-by: grYe99 --- .../model_executor/test_paddleocr_vl_mrope.py | 210 ++++++++++++++++++ vllm/model_executor/models/paddleocr_vl.py | 150 +++++-------- 2 files changed, 266 insertions(+), 94 deletions(-) create mode 100644 tests/model_executor/test_paddleocr_vl_mrope.py diff --git a/tests/model_executor/test_paddleocr_vl_mrope.py b/tests/model_executor/test_paddleocr_vl_mrope.py new file mode 100644 index 00000000000..9090a9254d0 --- /dev/null +++ b/tests/model_executor/test_paddleocr_vl_mrope.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass, field + +import pytest +import torch + +from vllm.model_executor.models.paddleocr_vl import ( + PaddleOCRVLForConditionalGeneration, +) +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalFieldElem, + MultiModalKwargsItem, + PlaceholderRange, +) + +pytestmark = pytest.mark.skip_global_cleanup + + +@pytest.fixture(autouse=True, scope="module") +def _force_cpu_default_device(): + original = torch.get_default_device() + torch.set_default_device("cpu") + yield + torch.set_default_device(original) + + +@dataclass +class DummyVisionConfig: + spatial_merge_size: int = 2 + patch_size: int = 14 + + +@dataclass +class DummyConfig: + image_token_id: int = 151655 + video_token_id: int = 151654 + vision_start_token_id: int = 151652 + vision_end_token_id: int = 151653 + vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig) + + +def make_model(config: DummyConfig) -> PaddleOCRVLForConditionalGeneration: + model = object.__new__(PaddleOCRVLForConditionalGeneration) + model.config = config + return model + + +def make_mm_feature( + *, + offset: int, + length: int, + image_grid_thw: tuple[int, int, int], +) -> MultiModalFeatureSpec: + return MultiModalFeatureSpec( + data=MultiModalKwargsItem( + { + "image_grid_thw": MultiModalFieldElem( + data=torch.tensor(image_grid_thw), + field=None, + ), + } + ), + modality="image", + identifier="DUMMY", + mm_position=PlaceholderRange(offset=offset, length=length), + ) + + +def test_get_mrope_input_positions_text_only(): + model = make_model(DummyConfig()) + input_tokens = [11, 12, 13, 14, 15] + positions, delta = model.get_mrope_input_positions( + input_tokens=input_tokens, + mm_features=[], + ) + expected = torch.tensor( + [ + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + ] + ) + assert torch.equal(positions, expected) + assert delta == 0 + + +def test_get_mrope_input_positions_single_image(): + model = make_model(DummyConfig()) + spatial_merge_size = model.config.vision_config.spatial_merge_size + + t, h, w = 1, 2, 2 + num_image_tokens = t * h * w + + input_tokens = ( + [10] + + [model.config.vision_start_token_id] + + [model.config.image_token_id] * num_image_tokens + + [model.config.vision_end_token_id] + + [30, 31] + ) + + mm_features = [ + make_mm_feature( + offset=2, # 1 (text) + 1 (vision_start) + length=num_image_tokens, + image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size), + ) + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=input_tokens, + mm_features=mm_features, + ) + + expected = torch.tensor( + [ + [0, 1, 2, 2, 2, 2, 4, 5, 6], + [0, 1, 2, 2, 3, 3, 4, 5, 6], + [0, 1, 2, 3, 2, 3, 4, 5, 6], + ] + ) + + assert torch.equal(positions, expected) + expected_delta = (positions.max().item() + 1) - len(input_tokens) + assert delta == expected_delta + + +def test_get_mrope_input_positions_multiple_images(): + model = make_model(DummyConfig()) + spatial_merge_size = model.config.vision_config.spatial_merge_size + + t1, h1, w1 = 1, 2, 2 + num1 = t1 * h1 * w1 + + t2, h2, w2 = 1, 1, 3 + num2 = t2 * h2 * w2 + + input_tokens = ( + [10] + + [model.config.vision_start_token_id] + + [model.config.image_token_id] * num1 + + [model.config.vision_end_token_id] + + [20, 21] + + [model.config.vision_start_token_id] + + [model.config.image_token_id] * num2 + + [model.config.vision_end_token_id] + + [30] + ) + + mm_features = [ + make_mm_feature( + offset=2, + length=num1, + image_grid_thw=(t1, h1 * spatial_merge_size, w1 * spatial_merge_size), + ), + make_mm_feature( + offset=2 + num1 + 1 + 2 + 1, + length=num2, + image_grid_thw=(t2, h2 * spatial_merge_size, w2 * spatial_merge_size), + ), + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=input_tokens, + mm_features=mm_features, + ) + + assert positions.shape == (3, 15) + assert not torch.equal(positions[:, 2:6], torch.arange(4).expand(3, 4) + 2) + assert not torch.equal(positions[:, 10:13], torch.arange(3).expand(3, 3) + 10) + + +def test_get_mrope_input_positions_image_at_start(): + model = make_model(DummyConfig()) + spatial_merge_size = model.config.vision_config.spatial_merge_size + + t, h, w = 1, 2, 2 + num_tokens = t * h * w + + input_tokens = ( + [model.config.vision_start_token_id] + + [model.config.image_token_id] * num_tokens + + [model.config.vision_end_token_id] + + [10, 11] + ) + + mm_features = [ + make_mm_feature( + offset=1, # start token at index 0 + length=num_tokens, + image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size), + ) + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=input_tokens, + mm_features=mm_features, + ) + + expected = torch.tensor( + [ + [0, 1, 1, 1, 1, 3, 4, 5], + [0, 1, 1, 2, 2, 3, 4, 5], + [0, 1, 2, 1, 2, 3, 4, 5], + ] + ) + + assert torch.equal(positions, expected) diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py index 48a285bc0f0..8eaf94620c1 100644 --- a/vllm/model_executor/models/paddleocr_vl.py +++ b/vllm/model_executor/models/paddleocr_vl.py @@ -15,7 +15,7 @@ # limitations under the License. import math -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from functools import partial from typing import Annotated, Literal @@ -1056,121 +1056,83 @@ class PaddleOCRVLForConditionalGeneration(nn.Module, SupportsMultiModal, Support ) -> torch.Tensor | None: return self.language_model.compute_logits(hidden_states) + def iter_mm_grid_thw( + self, mm_features: list[MultiModalFeatureSpec] + ) -> Iterator[tuple[int, int, int, int, float]]: + """ + Iterate over multimodal features and yield grid information. + + Args: + mm_features: List of multimodal feature specifications + + Yields: + Tuple of (offset, grid_t, grid_h, grid_w, t_factor) for each frame/image + """ + spatial_merge_size = self.config.vision_config.spatial_merge_size + tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0) + for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): + offset = mm_feature.mm_position.offset + if mm_feature.modality == "image": + t, h, w = mm_feature.data["image_grid_thw"].data.tolist() + assert t == 1, f"Image must have 1 frame, got {t}" + yield offset, 1, h // spatial_merge_size, w // spatial_merge_size, 1.0 + elif mm_feature.modality == "video": + t, h, w = mm_feature.data["video_grid_thw"].data.tolist() + second_per_grid_ts = 1.0 + if mm_feature.data.get("second_per_grid_ts", None): + second_per_grid_ts = mm_feature.data[ + "second_per_grid_ts" + ].data.item() + t_factor = second_per_grid_ts * tokens_per_second + yield ( + offset, + t, + h // spatial_merge_size, + w // spatial_merge_size, + t_factor, + ) + else: + raise ValueError(f"Unsupported modality: {mm_feature.modality}") + def get_mrope_input_positions( self, input_tokens: list[int], mm_features: list[MultiModalFeatureSpec], ) -> tuple[torch.Tensor, int]: - kwargs = MultiModalFeatureSpec.gather_kwargs( - mm_features, - {"image_grid_thw", "video_grid_thw", "second_per_grid_ts"}, - ) - image_grid_thw = [item.tolist() for item in kwargs.get("image_grid_thw", [])] - video_grid_thw = [item.tolist() for item in kwargs.get("video_grid_thw", [])] - second_per_grid_ts = kwargs.get("second_per_grid_ts", []) - - hf_config = self.config - image_token_id = hf_config.image_token_id - video_token_id = hf_config.video_token_id - vision_start_token_id = hf_config.vision_start_token_id - spatial_merge_size = hf_config.vision_config.spatial_merge_size - tokens_per_second = getattr(hf_config.vision_config, "tokens_per_second", 1.0) - - input_tokens_tensor = torch.tensor(input_tokens) - vision_start_indices = torch.argwhere( - input_tokens_tensor == vision_start_token_id - ).squeeze(1) - vision_tokens = input_tokens_tensor[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - - image_index, video_index = 0, 0 - for _ in range(image_nums + video_nums): - video_second_per_grid_t = 0.0 - if remain_images > 0: - try: - ed_image = input_tokens.index(image_token_id, st) - except ValueError: - ed_image = len(input_tokens) + 1 - else: - ed_image = len(input_tokens) + 1 - if remain_videos > 0: - try: - ed_video = input_tokens.index(video_token_id, st) - except ValueError: - ed_video = len(input_tokens) + 1 - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = image_grid_thw[image_index] - image_index += 1 - remain_images -= 1 - ed = ed_image - else: - t, h, w = video_grid_thw[video_index] - video_second_per_grid_t = 1.0 - if second_per_grid_ts: - video_second_per_grid_t = second_per_grid_ts[video_index] - video_index += 1 - remain_videos -= 1 - ed = ed_video - - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - text_len = ed - st + for ( + offset, + llm_grid_t, + llm_grid_h, + llm_grid_w, + t_factor, + ) in self.iter_mm_grid_thw(mm_features): + text_len = offset - st st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 llm_pos_ids_list.append( - torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx ) - t_index = ( - ( - torch.arange(llm_grid_t) - .view(-1, 1) - .expand(-1, llm_grid_h * llm_grid_w) - * video_second_per_grid_t - * tokens_per_second - ) - .long() - .flatten() - ) + grid_indices = np.indices((llm_grid_t, llm_grid_h, llm_grid_w)) + if t_factor != 1.0: + grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64) - h_index = ( - torch.arange(llm_grid_h) - .view(1, -1, 1) - .expand(llm_grid_t, -1, llm_grid_w) - .flatten() - ) - w_index = ( - torch.arange(llm_grid_w) - .view(1, 1, -1) - .expand(llm_grid_t, llm_grid_h, -1) - .flatten() - ) - llm_pos_ids_list.append( - torch.stack([t_index, h_index, w_index]) + text_len + st_idx - ) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w + llm_pos_ids_list.append(grid_indices.reshape(3, -1) + text_len + st_idx) + st = offset + llm_grid_t * llm_grid_h * llm_grid_w if st < len(input_tokens): st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 text_len = len(input_tokens) - st llm_pos_ids_list.append( - torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx ) - llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1) mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item() - return llm_positions, mrope_position_delta + return torch.from_numpy(llm_positions), mrope_position_delta def _parse_and_validate_image_input( self, **kwargs: object From 324a3d2bd8b5fc3f38c6d2f2cc243f747800ba28 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Thu, 16 Apr 2026 21:50:36 +0800 Subject: [PATCH 054/150] [CI/Build] Improve stability of CPU tests (#39966) Signed-off-by: jiang1.li --- .buildkite/hardware_tests/cpu.yaml | 4 ++-- tests/models/language/generation/test_common.py | 8 ++++---- tests/models/language/generation/test_granite.py | 1 + vllm/platforms/cpu.py | 1 + 4 files changed, 8 insertions(+), 6 deletions(-) diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index e466e2a5244..9b104444378 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -46,7 +46,7 @@ steps: - tests/models/language/pooling/ commands: - | - bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m " pytest -x -v -s tests/models/language/generation -m cpu_model pytest -x -v -s tests/models/language/pooling -m cpu_model" @@ -99,7 +99,7 @@ steps: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m " pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" - parallelism: 2 + parallelism: 3 - label: "Arm CPU Test" depends_on: [] diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py index b276f37a2a3..ebf847b706c 100644 --- a/tests/models/language/generation/test_common.py +++ b/tests/models/language/generation/test_common.py @@ -100,7 +100,7 @@ AITER_MODEL_LIST = [ pytest.param("bigcode/starcoder2-3b"), # starcoder2 pytest.param( "TitanML/tiny-mixtral", # mixtral - marks=[pytest.mark.core_model, pytest.mark.cpu_model], + marks=[pytest.mark.core_model], ), pytest.param("swiss-ai/Apertus-8B-Instruct-2509"), # apertus pytest.param( @@ -143,9 +143,9 @@ def test_models( # in parts of the operators pytest.skip(f"Skipping '{model}' model test with AITER kernel.") - if current_platform.is_cpu() and model == "TitanML/tiny-mixtral": - # This untrained model is sensitive to the rounding error - # Fuse ops to reduce bfloat16 rounding + if current_platform.is_cpu() and model in ("openai-community/gpt2",): + # These models are sensitive to the rounding error + # Fuse ops to reduce rounding monkeypatch.setenv("VLLM_CPU_CI_ENV", "0") with hf_runner(model) as hf_model: diff --git a/tests/models/language/generation/test_granite.py b/tests/models/language/generation/test_granite.py index e569e75ff3a..c0498b2f7de 100644 --- a/tests/models/language/generation/test_granite.py +++ b/tests/models/language/generation/test_granite.py @@ -15,6 +15,7 @@ MODELS = [ @pytest.mark.parametrize("dtype", ["bfloat16"]) @pytest.mark.parametrize("max_tokens", [64]) @pytest.mark.parametrize("num_logprobs", [5]) +@pytest.mark.cpu_model def test_models( hf_runner, vllm_runner, diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index f319dbc497a..86f07b1032e 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -242,6 +242,7 @@ class CpuPlatform(Platform): "cpp.dynamic_threads": True, } ) + compilation_config.ir_enable_torch_wrap = False if vllm_config.lora_config is not None: compilation_config.mode = CompilationMode.NONE From 5e5afafa21031abd7fb0e0707116e34be6f29af5 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 16 Apr 2026 10:52:58 -0400 Subject: [PATCH 055/150] [Doc] add docs for online quant frontend (#39736) Signed-off-by: Vasiliy Kuznetsov --- docs/features/quantization/README.md | 1 + docs/features/quantization/online.md | 94 ++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) create mode 100644 docs/features/quantization/online.md diff --git a/docs/features/quantization/README.md b/docs/features/quantization/README.md index 0b8fc71d3f3..4be088c56a4 100644 --- a/docs/features/quantization/README.md +++ b/docs/features/quantization/README.md @@ -16,6 +16,7 @@ The following are the supported quantization formats for vLLM: - [INT8 W8A8](int8.md) - [FP8 W8A8](fp8.md) - [NVIDIA Model Optimizer](modelopt.md) +- [Online Quantization](online.md) - [AMD Quark](quark.md) - [Quantized KV Cache](quantized_kvcache.md) - [TorchAO](torchao.md) diff --git a/docs/features/quantization/online.md b/docs/features/quantization/online.md new file mode 100644 index 00000000000..aa02366d2d5 --- /dev/null +++ b/docs/features/quantization/online.md @@ -0,0 +1,94 @@ +# Online Quantization + +Online quantization lets you take a BF16/FP16 model and quantize its Linear +and MoE weights to lower precision (such as FP8) at load time, without needing +a pre-quantized checkpoint or calibration data. Weights are converted during +model loading and activations are dynamically scaled during each forward pass. + +## Quick Start + +Pass a scheme name to the `quantization` parameter: + +```python +from vllm import LLM + +# Per-tensor FP8 quantization (one scale per weight tensor) +llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_tensor") + +# Per-block FP8 quantization (128x128 block scaling for weights and 1x128 block scaling for activations) +llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block") +``` + +Or with the CLI: + +```bash +vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor +vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block +``` + +## Supported Schemes + +| Scheme | Weight recipe | Activation recipe | Notes | +| ------ | ------------- | ------------------ | ----- | +| `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance | +| `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | | + +Support for additional schemes will be added in future versions of vllm. + +## Advanced Configuration + +For fine-grained control, use a `quantization_config` dictionary. + +### Separate Schemes for Dense and MoE Layers + +You can apply different quantization schemes to dense linear layers and MoE expert layers: + +```python +from vllm import LLM + +llm = LLM( + "ibm-granite/granite-3.0-1b-a400m-base", + quantization="fp8_per_tensor", + quantization_config={ + "linear_scheme_override": "fp8_per_block", + }, +) +``` + +Or, + +```python +from vllm import LLM + +llm = LLM( + "ibm-granite/granite-3.0-1b-a400m-base", + quantization="fp8_per_tensor", + quantization_config={ + "moe_scheme_override": "fp8_per_block", + }, +) +``` + +### Excluding Layers from Quantization + +Use the `ignore` parameter to skip specific layers. It accepts exact layer names and regex patterns (prefixed with `re:`): + +```python +from vllm import LLM + +llm = LLM( + "ibm-granite/granite-3.0-1b-a400m-base", + quantization="fp8_per_tensor", + quantization_config={ + "ignore": [ + # exact layer name + "model.layers.1.self_attn.o_proj", + # regex: skip all QKV projections + "re:.*[qkv]_proj", + ], + }, +) +``` + +!!! note + For fused layers (e.g., `qkv_proj` which fuses `q_proj`, `k_proj`, `v_proj`), the ignore pattern must match the **unfused** shard names (`q_proj`, `k_proj`, `v_proj`), not the fused name. From 4e8c3f1c197952454f244e5585a5de5990865939 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 16 Apr 2026 22:53:23 +0800 Subject: [PATCH 056/150] [Frontend][last/5] Improve pooling entrypoints | clean up. (#39675) Signed-off-by: wang.yuqi Signed-off-by: wang.yuqi Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- docs/models/pooling_models/README.md | 10 + docs/models/pooling_models/scoring.md | 4 + docs/models/pooling_models/token_classify.md | 2 +- docs/serving/openai_compatible_server.md | 19 +- vllm/entrypoints/llm.py | 2 +- vllm/entrypoints/openai/api_server.py | 30 +- vllm/entrypoints/openai/generate/factories.py | 42 +++ vllm/entrypoints/pooling/__init__.py | 130 --------- vllm/entrypoints/pooling/base/io_processor.py | 25 +- vllm/entrypoints/pooling/base/serving.py | 2 +- .../pooling/classify/api_router.py | 5 +- .../pooling/classify/io_processor.py | 2 +- vllm/entrypoints/pooling/classify/protocol.py | 9 +- vllm/entrypoints/pooling/classify/serving.py | 4 +- vllm/entrypoints/pooling/embed/api_router.py | 8 +- .../entrypoints/pooling/embed/io_processor.py | 39 +-- vllm/entrypoints/pooling/embed/protocol.py | 7 +- vllm/entrypoints/pooling/embed/serving.py | 27 +- vllm/entrypoints/pooling/factories.py | 256 ++++++++++++++++++ .../pooling/io_processor_factories.py | 76 ------ .../entrypoints/pooling/pooling/api_router.py | 5 +- .../pooling/pooling/io_processor.py | 4 +- vllm/entrypoints/pooling/pooling/protocol.py | 9 +- vllm/entrypoints/pooling/pooling/serving.py | 31 ++- .../pooling/scoring/io_processor.py | 12 +- vllm/entrypoints/pooling/scoring/protocol.py | 23 +- vllm/entrypoints/pooling/scoring/serving.py | 6 +- vllm/entrypoints/pooling/typing.py | 21 +- vllm/entrypoints/sagemaker/api_router.py | 84 +----- 29 files changed, 466 insertions(+), 428 deletions(-) create mode 100644 vllm/entrypoints/openai/generate/factories.py create mode 100644 vllm/entrypoints/pooling/factories.py delete mode 100644 vllm/entrypoints/pooling/io_processor_factories.py diff --git a/docs/models/pooling_models/README.md b/docs/models/pooling_models/README.md index 2cf721f5eef..838ec3bc901 100644 --- a/docs/models/pooling_models/README.md +++ b/docs/models/pooling_models/README.md @@ -59,6 +59,16 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models are a subset of classification models that accept two prompts as input and output num_labels equal to 1. +### Pooling Types + +| Pooling Tasks | Granularity | Description | +|----------------|---------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `CLS` pooling | Sequence-wise | For BERT‑like (bidirectional self‑attention) models, CLS pooling is used by default. This means the last_hidden_states corresponding to the first token (the [CLS] token) is taken as the output. | +| `LAST` pooling | Sequence-wise | For GPT‑like (causal self‑attention) models, LAST pooling is used by default. This means the last_hidden_states corresponding to the last token is taken as the output. | +| `MEAN` pooling | Sequence-wise | Many studies have shown that averaging the last_hidden_states over all input tokens performs better on certain downstream tasks. Therefore, more and more models are using MEAN pooling. | +| `ALL` pooling | Token-wise | Outputs the last_hidden_states for all input tokens. | +| `STEP` pooling | Token-wise | Filters and outputs the last_hidden_states corresponding to the token IDs returned by returned_token_ids. | + ### Score Types The scoring models is designed to compute similarity scores between two input prompts. It supports three model types diff --git a/docs/models/pooling_models/scoring.md b/docs/models/pooling_models/scoring.md index ac94a0cd76b..a4424642cd2 100644 --- a/docs/models/pooling_models/scoring.md +++ b/docs/models/pooling_models/scoring.md @@ -160,6 +160,8 @@ The following Score API parameters are supported: --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params" --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params" --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params" +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params" +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:score-request-params" ``` #### Examples @@ -370,6 +372,8 @@ The following rerank api parameters are supported: --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params" --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params" --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params" +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params" +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:rerank-request-params" ``` #### Examples diff --git a/docs/models/pooling_models/token_classify.md b/docs/models/pooling_models/token_classify.md index 201ce4ea6dc..945a710df96 100644 --- a/docs/models/pooling_models/token_classify.md +++ b/docs/models/pooling_models/token_classify.md @@ -68,7 +68,7 @@ If your model is not in the above list, we will try to automatically convert the Forced alignment usage requires `--hf-overrides '{"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}'`. Please refer to [examples/pooling/token_classify/forced_alignment_offline.py](../../../examples/pooling/token_classify/forced_alignment_offline.py). -### As Reward Models +### Reward Models Using token classification models as reward models. For details on reward models, see [Reward Models](reward.md). diff --git a/docs/serving/openai_compatible_server.md b/docs/serving/openai_compatible_server.md index 28e643ea5f0..cde7d597d82 100644 --- a/docs/serving/openai_compatible_server.md +++ b/docs/serving/openai_compatible_server.md @@ -467,28 +467,11 @@ It consists of two endpoints: - `/tokenize` corresponds to calling `tokenizer.encode()`. - `/detokenize` corresponds to calling `tokenizer.decode()`. -### Score API - -#### Score Template - -Some scoring models require a specific prompt format to work correctly. You can specify a custom score template using the `--chat-template` parameter (see [Chat Template](#chat-template)). - -Score templates are supported for **cross-encoder** models only. If you are using an **embedding** model for scoring, vLLM does not apply a score template. - -Like chat templates, the score template receives a `messages` list. For scoring, each message has a `role` attribute—either `"query"` or `"document"`. For the usual kind of point-wise cross-encoder, you can expect exactly two messages: one query and one document. To access the query and document content, use Jinja's `selectattr` filter: - -- **Query**: `{{ (messages | selectattr("role", "eq", "query") | first).content }}` -- **Document**: `{{ (messages | selectattr("role", "eq", "document") | first).content }}` - -This approach is more robust than index-based access (`messages[0]`, `messages[1]`) because it selects messages by their semantic role. It also avoids assumptions about message ordering if additional message types are added to `messages` in the future. - -Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../examples/pooling/score/template/nemotron-rerank.jinja) - ### Generative Scoring API The `/generative_scoring` endpoint uses a CausalLM model (e.g., Llama, Qwen, Mistral) to compute the probability of specified token IDs appearing as the next token. Each item (document) is concatenated with the query to form a prompt, and the model predicts how likely each label token is as the next token after that prompt. This lets you score items against a query — for example, asking "Is this the capital of France?" and scoring each city by how likely the model is to answer "Yes". -This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](#score-api), which uses cross-encoder, bi-encoder, or late-interaction models. +This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](../models/pooling_models/scoring.md#score-api), which uses cross-encoder, bi-encoder, or late-interaction models. **Requirements:** diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index d296e84d041..3a43dc591a2 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -48,7 +48,7 @@ from vllm.entrypoints.chat_utils import ( ChatTemplateContentFormatOption, load_chat_template, ) -from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors +from vllm.entrypoints.pooling.factories import init_pooling_io_processors from vllm.entrypoints.pooling.scoring.io_processor import ScoringIOProcessor from vllm.entrypoints.pooling.scoring.typing import ScoreInput from vllm.entrypoints.pooling.typing import OfflineInputsContext, OfflineOutputsContext diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 85d2fe43d01..9aac19e2fda 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -220,6 +220,12 @@ def build_app( elastic_ep_attach_router(app) + from vllm.entrypoints.openai.generative_scoring.api_router import ( + register_generative_scoring_api_router, + ) + + register_generative_scoring_api_router(app) + if "generate" in supported_tasks or "render" in supported_tasks: from vllm.entrypoints.serve.render.api_router import ( attach_router as attach_render_router, @@ -242,17 +248,10 @@ def build_app( register_realtime_api_router(app) if any(task in POOLING_TASKS for task in supported_tasks): - from vllm.entrypoints.pooling import register_pooling_api_routers + from vllm.entrypoints.pooling.factories import register_pooling_api_routers register_pooling_api_routers(app, supported_tasks, model_config) - if "generate" in supported_tasks: - from vllm.entrypoints.openai.generative_scoring.api_router import ( - register_generative_scoring_api_router, - ) - - register_generative_scoring_api_router(app) - app.root_path = args.root_path app.add_middleware( CORSMiddleware, @@ -401,6 +400,12 @@ async def init_app_state( engine_client, state, args, request_logger, supported_tasks ) + from vllm.entrypoints.openai.generative_scoring.api_router import ( + init_generative_scoring_state, + ) + + await init_generative_scoring_state(engine_client, state, args, request_logger) + if "transcription" in supported_tasks: from vllm.entrypoints.openai.speech_to_text.api_router import ( init_transcription_state, @@ -416,17 +421,10 @@ async def init_app_state( init_realtime_state(engine_client, state, args, request_logger, supported_tasks) if any(task in POOLING_TASKS for task in supported_tasks): - from vllm.entrypoints.pooling import init_pooling_state + from vllm.entrypoints.pooling.factories import init_pooling_state init_pooling_state(engine_client, state, args, request_logger, supported_tasks) - if "generate" in supported_tasks: - from vllm.entrypoints.openai.generative_scoring.api_router import ( - init_generative_scoring_state, - ) - - await init_generative_scoring_state(engine_client, state, args, request_logger) - state.enable_server_load_tracking = args.enable_server_load_tracking state.server_load_metrics = 0 diff --git a/vllm/entrypoints/openai/generate/factories.py b/vllm/entrypoints/openai/generate/factories.py new file mode 100644 index 00000000000..899601db3ca --- /dev/null +++ b/vllm/entrypoints/openai/generate/factories.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING + +from vllm.config import ModelConfig +from vllm.tasks import SupportedTask + +if TYPE_CHECKING: + from vllm.entrypoints.sagemaker.api_router import ( + EndpointFn, + GetHandlerFn, + RequestType, + ) + + +def get_generate_invocation_types( + supported_tasks: tuple["SupportedTask", ...], + model_config: ModelConfig | None = None, +): + # NOTE: Items defined earlier take higher priority + invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = [] + + if "generate" in supported_tasks: + from vllm.entrypoints.openai.chat_completion.api_router import ( + chat, + create_chat_completion, + ) + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ) + from vllm.entrypoints.openai.completion.api_router import ( + completion, + create_completion, + ) + from vllm.entrypoints.openai.completion.protocol import CompletionRequest + + invocation_types += [ + (ChatCompletionRequest, (chat, create_chat_completion)), + (CompletionRequest, (completion, create_completion)), + ] + + return invocation_types diff --git a/vllm/entrypoints/pooling/__init__.py b/vllm/entrypoints/pooling/__init__.py index 1980750ec20..e69de29bb2d 100644 --- a/vllm/entrypoints/pooling/__init__.py +++ b/vllm/entrypoints/pooling/__init__.py @@ -1,130 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from typing import TYPE_CHECKING - -from fastapi import FastAPI - -from vllm.config import ModelConfig -from vllm.entrypoints.pooling.utils import enable_scoring_api -from vllm.logger import init_logger - -if TYPE_CHECKING: - from argparse import Namespace - - from starlette.datastructures import State - - from vllm.engine.protocol import EngineClient - from vllm.entrypoints.logger import RequestLogger - from vllm.tasks import SupportedTask -else: - RequestLogger = object - SupportedTask = object - -logger = init_logger(__name__) - - -def register_pooling_api_routers( - app: FastAPI, - supported_tasks: tuple["SupportedTask", ...], - model_config: ModelConfig | None = None, -): - if model_config is None: - return - - pooling_task = model_config.get_pooling_task(supported_tasks) - - if pooling_task is not None: - from vllm.entrypoints.pooling.pooling.api_router import router as pooling_router - - app.include_router(pooling_router) - - if "classify" in supported_tasks: - from vllm.entrypoints.pooling.classify.api_router import ( - router as classify_router, - ) - - app.include_router(classify_router) - - if "embed" in supported_tasks: - from vllm.entrypoints.pooling.embed.api_router import router as embed_router - - app.include_router(embed_router) - - if enable_scoring_api(supported_tasks, model_config): - from vllm.entrypoints.pooling.scoring.api_router import router as score_router - - app.include_router(score_router) - - -def init_pooling_state( - engine_client: "EngineClient", - state: "State", - args: "Namespace", - request_logger: RequestLogger | None, - supported_tasks: tuple["SupportedTask", ...], -): - from vllm.entrypoints.chat_utils import load_chat_template - from vllm.entrypoints.pooling.classify.serving import ServingClassification - from vllm.entrypoints.pooling.embed.serving import ServingEmbedding - from vllm.entrypoints.pooling.pooling.serving import ServingPooling - from vllm.entrypoints.pooling.scoring.serving import ServingScores - from vllm.tasks import POOLING_TASKS - - model_config = engine_client.model_config - resolved_chat_template = load_chat_template(args.chat_template) - - state.serving_pooling = ( - ( - ServingPooling( - engine_client, - state.openai_serving_models, - supported_tasks=supported_tasks, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, - ) - ) - if any(t in supported_tasks for t in POOLING_TASKS) - else None - ) - state.serving_embedding = ( - ServingEmbedding( - engine_client, - state.openai_serving_models, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, - ) - if "embed" in supported_tasks - else None - ) - state.serving_classification = ( - ServingClassification( - engine_client, - state.openai_serving_models, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, - ) - if "classify" in supported_tasks - else None - ) - state.serving_scores = ( - ServingScores( - engine_client, - state.openai_serving_models, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, - enable_flash_late_interaction=getattr( - args, "enable_flash_late_interaction", True - ), - ) - if enable_scoring_api(supported_tasks, model_config) - else None - ) diff --git a/vllm/entrypoints/pooling/base/io_processor.py b/vllm/entrypoints/pooling/base/io_processor.py index 83e82664ef1..fc24bc65780 100644 --- a/vllm/entrypoints/pooling/base/io_processor.py +++ b/vllm/entrypoints/pooling/base/io_processor.py @@ -13,22 +13,29 @@ from vllm.entrypoints.chat_utils import ( ConversationMessage, ) from vllm.entrypoints.openai.engine.serving import RendererChatRequest, RendererRequest -from vllm.entrypoints.pooling.scoring.typing import ScoringData -from vllm.entrypoints.pooling.typing import ( - OfflineInputsContext, - OfflineOutputsContext, - PoolingChatLikeRequest, - PoolingCompletionLikeRequest, - PoolingServeContext, -) from vllm.inputs import EngineInput, SingletonPrompt from vllm.renderers import BaseRenderer, TokenizeParams, merge_kwargs from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq from vllm.tool_parsers import ToolParser from vllm.utils.mistral import is_mistral_tokenizer +from ..scoring.typing import ScoringData +from ..typing import ( + OfflineInputsContext, + OfflineOutputsContext, + PoolingChatLikeRequest, + PoolingCompletionLikeRequest, + PoolingServeContext, +) + class PoolingIOProcessor: + """Processor for handling preprocessing & postprocessing ops for pooling requests. + + This class manages both online (serving) and offline (batch) processing of pooling + requests, handling chat and completion formats. + """ + name: str def __init__( @@ -58,7 +65,7 @@ class PoolingIOProcessor: def pre_process_online(self, ctx: PoolingServeContext): request = ctx.request - if isinstance(ctx.request, PoolingChatLikeRequest): + if isinstance(request, PoolingChatLikeRequest): self._validate_chat_template( request_chat_template=request.chat_template, chat_template_kwargs=request.chat_template_kwargs, diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index c65bedc70f1..fe0b9fa8d95 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -22,7 +22,6 @@ from vllm.entrypoints.chat_utils import ( from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.pooling.typing import AnyPoolingRequest, PoolingServeContext from vllm.exceptions import VLLMNotFoundError from vllm.inputs import EngineInput from vllm.lora.request import LoRARequest @@ -36,6 +35,7 @@ from vllm.tracing import ( from vllm.utils import random_uuid from vllm.utils.async_utils import make_async, merge_async_iterators +from ..typing import AnyPoolingRequest, PoolingServeContext from .io_processor import PoolingIOProcessor diff --git a/vllm/entrypoints/pooling/classify/api_router.py b/vllm/entrypoints/pooling/classify/api_router.py index f254a6c2b39..2d27628bc69 100644 --- a/vllm/entrypoints/pooling/classify/api_router.py +++ b/vllm/entrypoints/pooling/classify/api_router.py @@ -5,13 +5,14 @@ from fastapi import APIRouter, Depends, Request from fastapi.responses import Response from vllm.entrypoints.openai.utils import validate_json_request -from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest -from vllm.entrypoints.pooling.classify.serving import ServingClassification from vllm.entrypoints.utils import ( load_aware_call, with_cancellation, ) +from .protocol import ClassificationRequest +from .serving import ServingClassification + router = APIRouter() diff --git a/vllm/entrypoints/pooling/classify/io_processor.py b/vllm/entrypoints/pooling/classify/io_processor.py index 9bb3774ab0b..1b2cd0ffa9d 100644 --- a/vllm/entrypoints/pooling/classify/io_processor.py +++ b/vllm/entrypoints/pooling/classify/io_processor.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor +from ..base.io_processor import PoolingIOProcessor class ClassifyIOProcessor(PoolingIOProcessor): diff --git a/vllm/entrypoints/pooling/classify/protocol.py b/vllm/entrypoints/pooling/classify/protocol.py index fe8c898e094..b5948e406be 100644 --- a/vllm/entrypoints/pooling/classify/protocol.py +++ b/vllm/entrypoints/pooling/classify/protocol.py @@ -9,15 +9,16 @@ from pydantic import Field from vllm import PoolingParams from vllm.config import ModelConfig from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo -from vllm.entrypoints.pooling.base.protocol import ( +from vllm.logger import init_logger +from vllm.renderers import TokenizeParams +from vllm.utils import random_uuid + +from ..base.protocol import ( ChatRequestMixin, ClassifyRequestMixin, CompletionRequestMixin, PoolingBasicRequestMixin, ) -from vllm.logger import init_logger -from vllm.renderers import TokenizeParams -from vllm.utils import random_uuid logger = init_logger(__name__) diff --git a/vllm/entrypoints/pooling/classify/serving.py b/vllm/entrypoints/pooling/classify/serving.py index a48ec819b7f..232483ca3ae 100644 --- a/vllm/entrypoints/pooling/classify/serving.py +++ b/vllm/entrypoints/pooling/classify/serving.py @@ -7,11 +7,11 @@ import numpy as np from fastapi.responses import JSONResponse from vllm.entrypoints.openai.engine.protocol import UsageInfo -from vllm.entrypoints.pooling.base.serving import PoolingServing -from vllm.entrypoints.pooling.typing import PoolingServeContext from vllm.logger import init_logger from vllm.outputs import ClassificationOutput +from ..base.serving import PoolingServing +from ..typing import PoolingServeContext from .io_processor import ClassifyIOProcessor from .protocol import ( ClassificationData, diff --git a/vllm/entrypoints/pooling/embed/api_router.py b/vllm/entrypoints/pooling/embed/api_router.py index 390efc6a13a..4eb86e4e2d2 100644 --- a/vllm/entrypoints/pooling/embed/api_router.py +++ b/vllm/entrypoints/pooling/embed/api_router.py @@ -7,13 +7,11 @@ from fastapi import APIRouter, Depends, Request from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.utils import validate_json_request -from vllm.entrypoints.pooling.embed.protocol import ( - CohereEmbedRequest, - EmbeddingRequest, -) -from vllm.entrypoints.pooling.embed.serving import ServingEmbedding from vllm.entrypoints.utils import load_aware_call, with_cancellation +from .protocol import CohereEmbedRequest, EmbeddingRequest +from .serving import ServingEmbedding + router = APIRouter() diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 09016253f09..8c28f9f3d4e 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -16,21 +16,6 @@ from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, CustomChatCompletionMessageParam, ) -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor -from vllm.entrypoints.pooling.embed.protocol import ( - CohereEmbedContent, - CohereEmbedInput, - CohereEmbedRequest, - EmbeddingChatRequest, - EmbeddingCompletionRequest, -) -from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessorMixin -from vllm.entrypoints.pooling.typing import ( - OfflineInputsContext, - PoolingChatLikeRequest, - PoolingCompletionLikeRequest, - PoolingServeContext, -) from vllm.inputs import EngineInput, tokens_input from vllm.logger import init_logger from vllm.outputs import PoolingOutput, PoolingRequestOutput @@ -39,6 +24,22 @@ from vllm.renderers.hf import resolve_chat_template from vllm.utils.collection_utils import chunk_list from vllm.utils.mistral import is_mistral_tokenizer +from ..base.io_processor import PoolingIOProcessor +from ..scoring.io_processor import JinaRankingIOProcessorMixin +from ..typing import ( + OfflineInputsContext, + PoolingChatLikeRequest, + PoolingCompletionLikeRequest, + PoolingServeContext, +) +from .protocol import ( + CohereEmbedContent, + CohereEmbedInput, + CohereEmbedRequest, + EmbeddingChatRequest, + EmbeddingCompletionRequest, +) + logger = init_logger(__name__) @@ -94,7 +95,7 @@ class EmbedIOProcessor(PoolingIOProcessor): if ctx.engine_inputs is None: raise ValueError("Engine prompts not available") - ctx.intermediates = ctx.engine_inputs + ctx.original_engine_inputs = ctx.engine_inputs request_id = ctx.request_id max_model_len = self.model_config.max_model_len chunked_engine_inputs: list[EngineInput] = [] @@ -189,10 +190,10 @@ class EmbedIOProcessor(PoolingIOProcessor): aggregator["total_weight"] += weight aggregator["chunk_count"] += 1 - if ctx.intermediates is None: - raise ValueError("Original prompts inputs not available") + if ctx.original_engine_inputs is None: + raise ValueError("Original engine inputs not available") - original_engine_inputs = cast(list[EngineInput], ctx.intermediates) + original_engine_inputs = ctx.original_engine_inputs num_prompts = len(original_engine_inputs) # Finalize aggregated results diff --git a/vllm/entrypoints/pooling/embed/protocol.py b/vllm/entrypoints/pooling/embed/protocol.py index 9b39b41df28..64087e2333e 100644 --- a/vllm/entrypoints/pooling/embed/protocol.py +++ b/vllm/entrypoints/pooling/embed/protocol.py @@ -18,14 +18,15 @@ from pydantic import BaseModel, Field from vllm import PoolingParams from vllm.config import ModelConfig from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo -from vllm.entrypoints.pooling.base.protocol import ( +from vllm.renderers import TokenizeParams +from vllm.utils import random_uuid + +from ..base.protocol import ( ChatRequestMixin, CompletionRequestMixin, EmbedRequestMixin, PoolingBasicRequestMixin, ) -from vllm.renderers import TokenizeParams -from vllm.utils import random_uuid # --------------------------------------------------------------------------- # OpenAI /v1/embeddings — request models diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index 9389309efc7..fd38598f7c7 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -9,9 +9,20 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse from typing_extensions import assert_never from vllm.entrypoints.openai.engine.protocol import UsageInfo -from vllm.entrypoints.pooling.base.serving import PoolingServing -from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor -from vllm.entrypoints.pooling.embed.protocol import ( +from vllm.logger import init_logger +from vllm.outputs import PoolingRequestOutput +from vllm.utils.serial_utils import EmbedDType, Endianness + +from ..base.serving import PoolingServing +from ..typing import PoolingServeContext +from ..utils import ( + encode_pooling_bytes, + encode_pooling_output_base64, + encode_pooling_output_float, + get_json_response_cls, +) +from .io_processor import EmbedIOProcessor +from .protocol import ( CohereBilledUnits, CohereEmbedRequest, CohereEmbedResponse, @@ -22,16 +33,6 @@ from vllm.entrypoints.pooling.embed.protocol import ( EmbeddingResponseData, build_typed_embeddings, ) -from vllm.entrypoints.pooling.typing import PoolingServeContext -from vllm.entrypoints.pooling.utils import ( - encode_pooling_bytes, - encode_pooling_output_base64, - encode_pooling_output_float, - get_json_response_cls, -) -from vllm.logger import init_logger -from vllm.outputs import PoolingRequestOutput -from vllm.utils.serial_utils import EmbedDType, Endianness logger = init_logger(__name__) diff --git a/vllm/entrypoints/pooling/factories.py b/vllm/entrypoints/pooling/factories.py new file mode 100644 index 00000000000..808004ed4d0 --- /dev/null +++ b/vllm/entrypoints/pooling/factories.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING + +from fastapi import FastAPI + +from vllm.config import ModelConfig, VllmConfig +from vllm.entrypoints.chat_utils import ChatTemplateConfig +from vllm.logger import init_logger +from vllm.plugins.io_processors import has_io_processor +from vllm.renderers import BaseRenderer +from vllm.tasks import POOLING_TASKS, SupportedTask + +from .base.io_processor import PoolingIOProcessor +from .utils import enable_scoring_api + +if TYPE_CHECKING: + from argparse import Namespace + + from starlette.datastructures import State + + from vllm.engine.protocol import EngineClient + from vllm.entrypoints.logger import RequestLogger + from vllm.entrypoints.sagemaker.api_router import ( + EndpointFn, + GetHandlerFn, + RequestType, + ) + +else: + RequestLogger = object + + +logger = init_logger(__name__) + + +def init_pooling_io_processors( + supported_tasks: tuple[SupportedTask, ...], + vllm_config: VllmConfig, + renderer: BaseRenderer, + chat_template_config: ChatTemplateConfig, +) -> dict[str, PoolingIOProcessor]: + model_config = vllm_config.model_config + processors: dict[str, type[PoolingIOProcessor]] = {} + + if "classify" in supported_tasks: + from .classify.io_processor import ClassifyIOProcessor + + processors["classify"] = ClassifyIOProcessor + + if "token_classify" in supported_tasks: + from .classify.io_processor import TokenClassifyIOProcessor + + processors["token_classify"] = TokenClassifyIOProcessor + + if "embed" in supported_tasks: + from .embed.io_processor import EmbedIOProcessor + + processors["embed"] = EmbedIOProcessor + + if "token_embed" in supported_tasks: + from .embed.io_processor import TokenEmbedIOProcessor + + processors["token_embed"] = TokenEmbedIOProcessor + + if has_io_processor( + vllm_config, + model_config.io_processor_plugin, + ): + from .pooling.io_processor import PluginWithIOProcessorPlugins + + processors["plugin"] = PluginWithIOProcessorPlugins + elif "plugin" in supported_tasks: + from .pooling.io_processor import PluginWithoutIOProcessorPlugins + + processors["plugin"] = PluginWithoutIOProcessorPlugins + + if enable_scoring_api(supported_tasks, model_config): + score_type = model_config.score_type + from .scoring.io_processor import ScoringIOProcessors + + if score_type is not None and score_type in ScoringIOProcessors: + processors[score_type] = ScoringIOProcessors[score_type] + + if model_config.architecture == "JinaForRanking": + from .embed.io_processor import JinaRankingTokenEmbedIOProcessor + from .scoring.io_processor import ScoringIOProcessors + + processors["token_embed"] = JinaRankingTokenEmbedIOProcessor + processors["late-interaction"] = ScoringIOProcessors["jina-reranking-scoring"] + + return { + task: processor_cls( + vllm_config=vllm_config, + renderer=renderer, + chat_template_config=chat_template_config, + ) + for task, processor_cls in processors.items() + } + + +def register_pooling_api_routers( + app: FastAPI, + supported_tasks: tuple["SupportedTask", ...], + model_config: ModelConfig | None = None, +): + if model_config is None: + return + + pooling_task = model_config.get_pooling_task(supported_tasks) + + if pooling_task is not None: + from .pooling.api_router import router as pooling_router + + app.include_router(pooling_router) + + if "classify" in supported_tasks: + from .classify.api_router import ( + router as classify_router, + ) + + app.include_router(classify_router) + + if "embed" in supported_tasks: + from .embed.api_router import router as embed_router + + app.include_router(embed_router) + + if enable_scoring_api(supported_tasks, model_config): + from .scoring.api_router import router as score_router + + app.include_router(score_router) + + +def init_pooling_state( + engine_client: "EngineClient", + state: "State", + args: "Namespace", + request_logger: RequestLogger | None, + supported_tasks: tuple["SupportedTask", ...], +): + from vllm.entrypoints.chat_utils import load_chat_template + from vllm.tasks import POOLING_TASKS + + from .classify.serving import ServingClassification + from .embed.serving import ServingEmbedding + from .pooling.serving import ServingPooling + from .scoring.serving import ServingScores + + model_config = engine_client.model_config + resolved_chat_template = load_chat_template(args.chat_template) + + state.serving_pooling = ( + ( + ServingPooling( + engine_client, + state.openai_serving_models, + supported_tasks=supported_tasks, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + ) + ) + if any(t in supported_tasks for t in POOLING_TASKS) + else None + ) + state.serving_embedding = ( + ServingEmbedding( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + ) + if "embed" in supported_tasks + else None + ) + state.serving_classification = ( + ServingClassification( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + ) + if "classify" in supported_tasks + else None + ) + state.serving_scores = ( + ServingScores( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + enable_flash_late_interaction=getattr( + args, "enable_flash_late_interaction", True + ), + ) + if enable_scoring_api(supported_tasks, model_config) + else None + ) + + +def get_pooling_invocation_types( + supported_tasks: tuple["SupportedTask", ...], + model_config: ModelConfig | None = None, +): + # NOTE: Items defined earlier take higher priority + invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = [] + + if "embed" in supported_tasks: + from .embed.api_router import create_embedding, embedding + from .embed.protocol import EmbeddingRequest + + invocation_types += [ + (EmbeddingRequest, (embedding, create_embedding)), + ] + + if "classify" in supported_tasks: + from .classify.api_router import classify, create_classify + from .classify.protocol import ClassificationRequest + + invocation_types += [ + (ClassificationRequest, (classify, create_classify)), + ] + + if enable_scoring_api(supported_tasks, model_config): + from .scoring.api_router import do_rerank, rerank + from .scoring.protocol import RerankRequest + + invocation_types += [ + (RerankRequest, (rerank, do_rerank)), + ] + + from .scoring.api_router import create_score, score + from .scoring.protocol import ScoreRequest + + invocation_types += [ + (ScoreRequest, (score, create_score)), + ] + + if any(task in POOLING_TASKS for task in supported_tasks): + from .pooling.api_router import create_pooling, pooling + from .pooling.protocol import PoolingRequest + + invocation_types += [ + (PoolingRequest, (pooling, create_pooling)), + ] + + return invocation_types diff --git a/vllm/entrypoints/pooling/io_processor_factories.py b/vllm/entrypoints/pooling/io_processor_factories.py deleted file mode 100644 index 5e67d069da7..00000000000 --- a/vllm/entrypoints/pooling/io_processor_factories.py +++ /dev/null @@ -1,76 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ChatTemplateConfig -from vllm.plugins.io_processors import has_io_processor -from vllm.renderers import BaseRenderer -from vllm.tasks import SupportedTask - -from .base.io_processor import PoolingIOProcessor -from .utils import enable_scoring_api - - -def init_pooling_io_processors( - supported_tasks: tuple[SupportedTask, ...], - vllm_config: VllmConfig, - renderer: BaseRenderer, - chat_template_config: ChatTemplateConfig, -) -> dict[str, PoolingIOProcessor]: - model_config = vllm_config.model_config - processors: dict[str, type[PoolingIOProcessor]] = {} - - if "classify" in supported_tasks: - from .classify.io_processor import ClassifyIOProcessor - - processors["classify"] = ClassifyIOProcessor - - if "token_classify" in supported_tasks: - from .classify.io_processor import TokenClassifyIOProcessor - - processors["token_classify"] = TokenClassifyIOProcessor - - if "embed" in supported_tasks: - from .embed.io_processor import EmbedIOProcessor - - processors["embed"] = EmbedIOProcessor - - if "token_embed" in supported_tasks: - from .embed.io_processor import TokenEmbedIOProcessor - - processors["token_embed"] = TokenEmbedIOProcessor - - if has_io_processor( - vllm_config, - model_config.io_processor_plugin, - ): - from .pooling.io_processor import PluginWithIOProcessorPlugins - - processors["plugin"] = PluginWithIOProcessorPlugins - elif "plugin" in supported_tasks: - from .pooling.io_processor import PluginWithoutIOProcessorPlugins - - processors["plugin"] = PluginWithoutIOProcessorPlugins - - if enable_scoring_api(supported_tasks, model_config): - score_type = model_config.score_type - from .scoring.io_processor import ScoringIOProcessors - - if score_type is not None and score_type in ScoringIOProcessors: - processors[score_type] = ScoringIOProcessors[score_type] - - if model_config.architecture == "JinaForRanking": - from .embed.io_processor import JinaRankingTokenEmbedIOProcessor - from .scoring.io_processor import ScoringIOProcessors - - processors["token_embed"] = JinaRankingTokenEmbedIOProcessor - processors["late-interaction"] = ScoringIOProcessors["jina-reranking-scoring"] - - return { - task: processor_cls( - vllm_config=vllm_config, - renderer=renderer, - chat_template_config=chat_template_config, - ) - for task, processor_cls in processors.items() - } diff --git a/vllm/entrypoints/pooling/pooling/api_router.py b/vllm/entrypoints/pooling/pooling/api_router.py index a08570038c3..0c77c050dc0 100644 --- a/vllm/entrypoints/pooling/pooling/api_router.py +++ b/vllm/entrypoints/pooling/pooling/api_router.py @@ -6,10 +6,11 @@ from fastapi import APIRouter, Depends, Request from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.utils import validate_json_request -from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest -from vllm.entrypoints.pooling.pooling.serving import ServingPooling from vllm.entrypoints.utils import load_aware_call, with_cancellation +from .protocol import PoolingRequest +from .serving import ServingPooling + router = APIRouter() diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index 31f860144ea..b07ceeede32 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -4,12 +4,12 @@ from collections.abc import Sequence from typing import Any from vllm import PoolingParams, PoolingRequestOutput -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor from vllm.inputs import EngineInput from vllm.logger import init_logger from vllm.plugins.io_processors import get_io_processor from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq +from ..base.io_processor import PoolingIOProcessor from ..typing import OfflineInputsContext, OfflineOutputsContext, PoolingServeContext from .protocol import IOProcessorRequest, IOProcessorResponse @@ -17,6 +17,8 @@ logger = init_logger(__name__) class PluginWithoutIOProcessorPlugins(PoolingIOProcessor): + # Some models, such as Terratorch (tests/models/test_terratorch.py), + # use plugin tasks in the pooler but do not use IO Processor plugins. name = "plugin" diff --git a/vllm/entrypoints/pooling/pooling/protocol.py b/vllm/entrypoints/pooling/pooling/protocol.py index 098690db262..b75c34e9790 100644 --- a/vllm/entrypoints/pooling/pooling/protocol.py +++ b/vllm/entrypoints/pooling/pooling/protocol.py @@ -8,7 +8,11 @@ from pydantic import Field from vllm import PoolingParams from vllm.config import ModelConfig from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo -from vllm.entrypoints.pooling.base.protocol import ( +from vllm.renderers import TokenizeParams +from vllm.tasks import PoolingTask +from vllm.utils import random_uuid + +from ..base.protocol import ( ChatRequestMixin, ClassifyRequestMixin, CompletionRequestMixin, @@ -16,9 +20,6 @@ from vllm.entrypoints.pooling.base.protocol import ( EncodingRequestMixin, PoolingBasicRequestMixin, ) -from vllm.renderers import TokenizeParams -from vllm.tasks import PoolingTask -from vllm.utils import random_uuid class PoolingCompletionRequest( diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index ea8eff6eb30..3498e830665 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -11,27 +11,28 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse from typing_extensions import assert_never from vllm.entrypoints.openai.engine.protocol import UsageInfo -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor -from vllm.entrypoints.pooling.base.serving import PoolingServingBase -from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors -from vllm.entrypoints.pooling.pooling.protocol import ( +from vllm.logger import init_logger +from vllm.outputs import PoolingRequestOutput +from vllm.tasks import SupportedTask +from vllm.utils.serial_utils import EmbedDType, Endianness + +from ..base.io_processor import PoolingIOProcessor +from ..base.serving import PoolingServingBase +from ..factories import init_pooling_io_processors +from ..typing import AnyPoolingRequest, PoolingServeContext +from ..utils import ( + encode_pooling_bytes, + encode_pooling_output_base64, + encode_pooling_output_float, + get_json_response_cls, +) +from .protocol import ( IOProcessorRequest, PoolingBytesResponse, PoolingRequest, PoolingResponse, PoolingResponseData, ) -from vllm.entrypoints.pooling.typing import AnyPoolingRequest, PoolingServeContext -from vllm.entrypoints.pooling.utils import ( - encode_pooling_bytes, - encode_pooling_output_base64, - encode_pooling_output_float, - get_json_response_cls, -) -from vllm.logger import init_logger -from vllm.outputs import PoolingRequestOutput -from vllm.tasks import SupportedTask -from vllm.utils.serial_utils import EmbedDType, Endianness logger = init_logger(__name__) diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 549bae2775d..e706601c568 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -7,12 +7,6 @@ from typing import Any, TypeAlias import torch.nn.functional as F from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor -from vllm.entrypoints.pooling.typing import ( - OfflineInputsContext, - OfflineOutputsContext, - PoolingServeContext, -) from vllm.inputs import EngineInput from vllm.renderers import TokenizeParams from vllm.renderers.hf import safe_apply_chat_template @@ -20,6 +14,12 @@ from vllm.tasks import PoolingTask from vllm.utils.mistral import is_mistral_tokenizer from ...chat_utils import ChatTemplateResolutionError +from ..base.io_processor import PoolingIOProcessor +from ..typing import ( + OfflineInputsContext, + OfflineOutputsContext, + PoolingServeContext, +) from .protocol import RerankRequest, ScoreRequest, ScoringRequest from .typing import ScoreData, ScoreInput, ScoringData from .utils import ( diff --git a/vllm/entrypoints/pooling/scoring/protocol.py b/vllm/entrypoints/pooling/scoring/protocol.py index 83fdafb1458..d32a3660d02 100644 --- a/vllm/entrypoints/pooling/scoring/protocol.py +++ b/vllm/entrypoints/pooling/scoring/protocol.py @@ -8,18 +8,16 @@ from pydantic import BaseModel, Field from vllm import PoolingParams from vllm.config import ModelConfig from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo -from vllm.entrypoints.pooling.base.protocol import ( - ClassifyRequestMixin, - PoolingBasicRequestMixin, -) from vllm.renderers import TokenizeParams from vllm.tasks import PoolingTask from vllm.utils import random_uuid +from ..base.protocol import ClassifyRequestMixin, PoolingBasicRequestMixin from .typing import ScoreContentPartParam, ScoreInput -class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin): +class ScoringRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin): + # --8<-- [start:scoring-common-params] max_tokens_per_query: int = Field( default=0, description=( @@ -37,6 +35,7 @@ class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin): "applies to the combined query+document)." ), ) + # --8<-- [end:scoring-common-params] def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: encoder_config = model_config.encoder_config or {} @@ -57,14 +56,16 @@ class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin): ) -class ScoreDataRequest(ScoreRequestMixin): +class ScoreDataRequest(ScoringRequestMixin): data_1: ScoreInput | list[ScoreInput] data_2: ScoreInput | list[ScoreInput] -class ScoreQueriesDocumentsRequest(ScoreRequestMixin): +class ScoreQueriesDocumentsRequest(ScoringRequestMixin): + # --8<-- [start:score-request-params] queries: ScoreInput | list[ScoreInput] documents: ScoreInput | list[ScoreInput] + # --8<-- [end:score-request-params] @property def data_1(self): @@ -75,7 +76,7 @@ class ScoreQueriesDocumentsRequest(ScoreRequestMixin): return self.documents -class ScoreQueriesItemsRequest(ScoreRequestMixin): +class ScoreQueriesItemsRequest(ScoringRequestMixin): queries: ScoreInput | list[ScoreInput] items: ScoreInput | list[ScoreInput] @@ -88,7 +89,7 @@ class ScoreQueriesItemsRequest(ScoreRequestMixin): return self.items -class ScoreTextRequest(ScoreRequestMixin): +class ScoreTextRequest(ScoringRequestMixin): text_1: ScoreInput | list[ScoreInput] text_2: ScoreInput | list[ScoreInput] @@ -109,10 +110,12 @@ ScoreRequest: TypeAlias = ( ) -class RerankRequest(ScoreRequestMixin): +class RerankRequest(ScoringRequestMixin): + # --8<-- [start:rerank-request-params] query: ScoreInput documents: ScoreInput | list[ScoreInput] top_n: int = Field(default_factory=lambda: 0) + # --8<-- [end:rerank-request-params] ScoringRequest: TypeAlias = ScoreRequest | RerankRequest diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py index df866efd56e..dc0d7ccd2bf 100644 --- a/vllm/entrypoints/pooling/scoring/serving.py +++ b/vllm/entrypoints/pooling/scoring/serving.py @@ -6,8 +6,6 @@ from fastapi.responses import JSONResponse, Response from vllm import PoolingParams from vllm.engine.protocol import EngineClient from vllm.entrypoints.openai.engine.protocol import UsageInfo -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor -from vllm.entrypoints.pooling.base.serving import PoolingServing from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput from vllm.v1.pool.late_interaction import ( @@ -15,6 +13,8 @@ from vllm.v1.pool.late_interaction import ( build_late_interaction_query_params, ) +from ..base.io_processor import PoolingIOProcessor +from ..base.serving import PoolingServing from .io_processor import ScoringIOProcessors, ScoringServeContext from .protocol import ( RerankDocument, @@ -90,7 +90,7 @@ class ServingScores(PoolingServing): ctx.request.top_n if ctx.request.top_n > 0 else len(final_res_batch), ) else: - raise NotImplementedError("") + raise ValueError(f"Invalid {self.request_id_prefix} request type") def _request_output_to_score_response( self, diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 0ce9d5c8384..ffcd3e7be43 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -9,29 +9,30 @@ from fastapi import Request from pydantic import ConfigDict from vllm import PoolingParams, PoolingRequestOutput, PromptType -from vllm.entrypoints.pooling.classify.protocol import ( +from vllm.inputs import DataPrompt, EngineInput +from vllm.lora.request import LoRARequest + +from .classify.protocol import ( ClassificationChatRequest, ClassificationCompletionRequest, ClassificationResponse, ) -from vllm.entrypoints.pooling.embed.protocol import ( +from .embed.protocol import ( CohereEmbedRequest, EmbeddingBytesResponse, EmbeddingChatRequest, EmbeddingCompletionRequest, EmbeddingResponse, ) -from vllm.entrypoints.pooling.pooling.protocol import ( +from .pooling.protocol import ( IOProcessorRequest, PoolingBytesResponse, PoolingChatRequest, PoolingCompletionRequest, PoolingResponse, ) -from vllm.entrypoints.pooling.scoring.protocol import ScoringRequest, ScoringResponse -from vllm.entrypoints.pooling.scoring.typing import ScoringData -from vllm.inputs import DataPrompt, EngineInput -from vllm.lora.request import LoRARequest +from .scoring.protocol import ScoringRequest, ScoringResponse +from .scoring.typing import ScoringData PoolingCompletionLikeRequest: TypeAlias = ( EmbeddingCompletionRequest @@ -65,6 +66,8 @@ PoolingRequestT = TypeVar("PoolingRequestT", bound=AnyPoolingRequest) @dataclass(kw_only=True) class PoolingServeContext(Generic[PoolingRequestT]): + model_config = ConfigDict(arbitrary_types_allowed=True) + request: PoolingRequestT raw_request: Request | None = None model_name: str @@ -74,14 +77,14 @@ class PoolingServeContext(Generic[PoolingRequestT]): lora_request: LoRARequest | None = None engine_inputs: Sequence[EngineInput] | None = None prompt_request_ids: list[str] | None = None - intermediates: Any | None = None result_generator: AsyncGenerator[tuple[int, PoolingRequestOutput], None] | None = ( None ) final_res_batch: list[PoolingRequestOutput] = field(default_factory=list) - model_config = ConfigDict(arbitrary_types_allowed=True) + ## for Long Text Embedding with Chunked Processing + original_engine_inputs: Sequence[EngineInput] | None = None ## for bi-encoder & late-interaction n_queries: int | None = None diff --git a/vllm/entrypoints/sagemaker/api_router.py b/vllm/entrypoints/sagemaker/api_router.py index 1d63793df39..b3b11cd07b4 100644 --- a/vllm/entrypoints/sagemaker/api_router.py +++ b/vllm/entrypoints/sagemaker/api_router.py @@ -13,12 +13,13 @@ from fastapi.responses import JSONResponse, Response from vllm.config import ModelConfig from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.engine.serving import OpenAIServing +from vllm.entrypoints.openai.generate.factories import get_generate_invocation_types from vllm.entrypoints.openai.utils import validate_json_request from vllm.entrypoints.pooling.base.serving import PoolingServingBase -from vllm.entrypoints.pooling.utils import enable_scoring_api +from vllm.entrypoints.pooling.factories import get_pooling_invocation_types from vllm.entrypoints.serve.instrumentator.basic import base from vllm.entrypoints.serve.instrumentator.health import health -from vllm.tasks import POOLING_TASKS, SupportedTask +from vllm.tasks import SupportedTask # TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers # (requires typing_extensions >= 4.13) @@ -27,80 +28,6 @@ GetHandlerFn = Callable[[Request], OpenAIServing | PoolingServingBase | None] EndpointFn = Callable[[RequestType, Request], Awaitable[Any]] -def get_invocation_types( - supported_tasks: tuple["SupportedTask", ...], - model_config: ModelConfig | None = None, -): - # NOTE: Items defined earlier take higher priority - INVOCATION_TYPES: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = [] - - if "generate" in supported_tasks: - from vllm.entrypoints.openai.chat_completion.api_router import ( - chat, - create_chat_completion, - ) - from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, - ) - from vllm.entrypoints.openai.completion.api_router import ( - completion, - create_completion, - ) - from vllm.entrypoints.openai.completion.protocol import CompletionRequest - - INVOCATION_TYPES += [ - (ChatCompletionRequest, (chat, create_chat_completion)), - (CompletionRequest, (completion, create_completion)), - ] - - if "embed" in supported_tasks: - from vllm.entrypoints.pooling.embed.api_router import ( - create_embedding, - embedding, - ) - from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest - - INVOCATION_TYPES += [ - (EmbeddingRequest, (embedding, create_embedding)), - ] - - if "classify" in supported_tasks: - from vllm.entrypoints.pooling.classify.api_router import ( - classify, - create_classify, - ) - from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest - - INVOCATION_TYPES += [ - (ClassificationRequest, (classify, create_classify)), - ] - - if enable_scoring_api(supported_tasks, model_config): - from vllm.entrypoints.pooling.scoring.api_router import do_rerank, rerank - from vllm.entrypoints.pooling.scoring.protocol import RerankRequest - - INVOCATION_TYPES += [ - (RerankRequest, (rerank, do_rerank)), - ] - - from vllm.entrypoints.pooling.scoring.api_router import create_score, score - from vllm.entrypoints.pooling.scoring.protocol import ScoreRequest - - INVOCATION_TYPES += [ - (ScoreRequest, (score, create_score)), - ] - - if any(task in POOLING_TASKS for task in supported_tasks): - from vllm.entrypoints.pooling.pooling.api_router import create_pooling, pooling - from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest - - INVOCATION_TYPES += [ - (PoolingRequest, (pooling, create_pooling)), - ] - - return INVOCATION_TYPES - - def attach_router( app: FastAPI, supported_tasks: tuple["SupportedTask", ...], @@ -109,7 +36,10 @@ def attach_router( router = APIRouter() # NOTE: Construct the TypeAdapters only once - INVOCATION_TYPES = get_invocation_types(supported_tasks, model_config) + INVOCATION_TYPES = get_generate_invocation_types( + supported_tasks, model_config + ) + get_pooling_invocation_types(supported_tasks, model_config) + INVOCATION_VALIDATORS = [ (pydantic.TypeAdapter(request_type), (get_handler, endpoint)) for request_type, (get_handler, endpoint) in INVOCATION_TYPES From a302a8fd1b2033fd3ea3981556af525174b6ff16 Mon Sep 17 00:00:00 2001 From: daiyu1111 <145559357+daiyu1111@users.noreply.github.com> Date: Thu, 16 Apr 2026 22:56:06 +0800 Subject: [PATCH 057/150] [Bugfix] Fix LLM priority normalization for single-string prompts (#40011) Signed-off-by: daiyu1111 <2356690121@qq.com> --- tests/entrypoints/llm/test_generate.py | 6 ++++++ vllm/entrypoints/llm.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/llm/test_generate.py b/tests/entrypoints/llm/test_generate.py index 34465b7d270..82f38adfb77 100644 --- a/tests/entrypoints/llm/test_generate.py +++ b/tests/entrypoints/llm/test_generate.py @@ -91,6 +91,12 @@ def test_multiple_priority(llm: LLM): outputs = llm.generate(PROMPTS, sampling_params=None, priority=[]) +def test_single_prompt_priority(llm: LLM): + # Single string prompts should be normalized to one request. + outputs = llm.generate(PROMPTS[0], sampling_params=None, priority=[0]) + assert len(outputs) == 1 + + def test_max_model_len(): max_model_len = 20 llm = LLM( diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 3a43dc591a2..d135a9ec976 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1575,7 +1575,7 @@ class LLM: seq_prompts = prompt_to_seq(prompts) seq_params = self._params_to_seq(params, len(seq_prompts)) seq_lora_requests = self._lora_request_to_seq(lora_request, len(seq_prompts)) - seq_priority = self._priority_to_seq(priority, len(prompts)) + seq_priority = self._priority_to_seq(priority, len(seq_prompts)) return self._render_and_add_requests( prompts=( From 3daca38e2279538b420641bd41853c19e5ad01e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 16 Apr 2026 17:08:22 +0200 Subject: [PATCH 058/150] [Misc] `toy_proxy_server` handle min_tokens (#39706) Signed-off-by: NickLucche --- tests/v1/kv_connector/nixl_integration/toy_proxy_server.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py b/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py index b92d3fcd6fb..1a910c8d24d 100644 --- a/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py +++ b/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py @@ -173,6 +173,9 @@ async def send_request_to_service( req_data["max_completion_tokens"] = 1 if "stream_options" in req_data: del req_data["stream_options"] + # These args are not supported for P + min_tokens = req_data.pop("min_tokens", None) + min_completion_tokens = req_data.pop("min_completion_tokens", None) headers = { "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", "X-Request-Id": request_id, @@ -187,6 +190,10 @@ async def send_request_to_service( # otherwise, it would http.ReadError await response.aread() + # Add back the min_tokens and min_completion_tokens so D can use them + req_data["min_tokens"] = min_tokens + req_data["min_completion_tokens"] = min_completion_tokens + return response From 82531edbfb6b33e1c8667dea15c8622f011dcef0 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Thu, 16 Apr 2026 23:48:17 +0800 Subject: [PATCH 059/150] [Refactor] Remove `resampy` dependency (#39524) Signed-off-by: Isotr0py --- requirements/test/cuda.in | 1 - requirements/test/cuda.txt | 4 ---- requirements/test/rocm.in | 1 - requirements/test/rocm.txt | 4 ---- setup.py | 1 - vllm/multimodal/audio.py | 20 +------------------- vllm/multimodal/media/audio.py | 9 ++------- 7 files changed, 3 insertions(+), 37 deletions(-) diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 5cf3a69e1fb..30076ff7071 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -21,7 +21,6 @@ vocos # required for minicpmo_26 test peft>=0.18.1 # required for phi-4-mm test pqdm ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests -resampy # required for audio tests sentence-transformers>=5.2.0 # required for embedding tests soundfile # required for audio tests jiwer # required for audio tests diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index ed67685e6eb..449939b965c 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -555,7 +555,6 @@ numba==0.61.2 # -c requirements/cuda.txt # -r requirements/test/cuda.in # librosa - # resampy numpy==2.2.6 # via # -r requirements/test/cuda.in @@ -596,7 +595,6 @@ numpy==2.2.6 # pyogrio # pywavelets # rasterio - # resampy # rioxarray # rouge-score # runai-model-streamer @@ -1015,8 +1013,6 @@ requests==2.32.3 # tacoreader # tiktoken # wandb -resampy==0.4.3 - # via -r requirements/test/cuda.in responses==0.25.3 # via genai-perf rfc3339-validator==0.1.4 diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index dbb1500edcf..0c9335240f5 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -23,7 +23,6 @@ vocos # required for minicpmo_26 test peft>=0.15.0 # required for phi-4-mm test pqdm ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests -resampy # required for audio tests sentence-transformers>=5.2.0 # required for embedding tests soundfile # required for audio tests jiwer # required for audio tests diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index ba9cd3dfdcf..2e03df1bdc0 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -662,7 +662,6 @@ numba==0.61.2 # -c requirements/rocm.txt # -r requirements/test/rocm.in # librosa - # resampy numkong==7.1.1 # via albucore numpy==2.2.6 @@ -708,7 +707,6 @@ numpy==2.2.6 # pytrec-eval-terrier # pywavelets # rasterio - # resampy # rioxarray # rouge-score # runai-model-streamer @@ -1193,8 +1191,6 @@ requests==2.32.5 # tacoreader # tiktoken # wandb -resampy==0.4.3 - # via -r requirements/test/rocm.in responses==0.26.0 # via genai-perf rfc3339-validator==0.1.4 diff --git a/setup.py b/setup.py index af52cb58dfc..37782062ae0 100644 --- a/setup.py +++ b/setup.py @@ -1093,7 +1093,6 @@ setup( "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"], "audio": [ "av", - "resampy", "scipy", "soundfile", "mistral_common[audio]", diff --git a/vllm/multimodal/audio.py b/vllm/multimodal/audio.py index 0a748a6d15c..c232bb9ea6a 100644 --- a/vllm/multimodal/audio.py +++ b/vllm/multimodal/audio.py @@ -16,11 +16,6 @@ try: except ImportError: av = PlaceholderModule("av") # type: ignore[assignment] -try: - import resampy -except ImportError: - resampy = PlaceholderModule("resampy") # type: ignore[assignment] - try: import scipy.signal as scipy_signal except ImportError: @@ -229,15 +224,6 @@ def resample_audio_pyav( return result[:expected_len] -def resample_audio_resampy( - audio: npt.NDArray[np.floating], - *, - orig_sr: float, - target_sr: float, -) -> npt.NDArray[np.floating]: - return resampy.resample(audio, sr_orig=orig_sr, sr_new=target_sr) - - def resample_audio_scipy( audio: npt.NDArray[np.floating], *, @@ -257,7 +243,7 @@ class AudioResampler: def __init__( self, target_sr: float | None = None, - method: Literal["pyav", "resampy", "scipy"] = "resampy", + method: Literal["pyav", "scipy"] = "pyav", ): self.target_sr = target_sr self.method = method @@ -281,10 +267,6 @@ class AudioResampler: return audio if self.method == "pyav": return resample_audio_pyav(audio, orig_sr=orig_sr, target_sr=self.target_sr) - if self.method == "resampy": - return resample_audio_resampy( - audio, orig_sr=orig_sr, target_sr=self.target_sr - ) elif self.method == "scipy": return resample_audio_scipy( audio, orig_sr=orig_sr, target_sr=self.target_sr diff --git a/vllm/multimodal/media/audio.py b/vllm/multimodal/media/audio.py index 4d41248f147..26d58d13a73 100644 --- a/vllm/multimodal/media/audio.py +++ b/vllm/multimodal/media/audio.py @@ -10,6 +10,7 @@ import pybase64 import torch from vllm.logger import init_logger +from vllm.multimodal.audio import resample_audio_pyav from vllm.utils.import_utils import PlaceholderModule from vllm.utils.serial_utils import tensor2base64 @@ -28,12 +29,6 @@ except ImportError: soundfile = PlaceholderModule("soundfile") # type: ignore[assignment] -try: - import resampy -except ImportError: - resampy = PlaceholderModule("resampy") # type: ignore[assignment] - - # Public libsndfile error codes exposed via `soundfile.LibsndfileError.code`, soundfile # being librosa's main backend. Used to validate if an audio loading error is due to a # server error vs a client error (invalid audio file). @@ -129,7 +124,7 @@ def load_audio_soundfile( y = np.mean(y, axis=tuple(range(y.ndim - 1))) if sr is not None and sr != native_sr: - y = resampy.resample(y, sr_orig=native_sr, sr_new=sr) + y = resample_audio_pyav(y, orig_sr=native_sr, target_sr=sr) return y, int(sr) return y, native_sr From 692db29cd45fcd2d2e0c3b7259b608ec1f5855ef Mon Sep 17 00:00:00 2001 From: Nikita Shapovalov Date: Thu, 16 Apr 2026 17:49:29 +0200 Subject: [PATCH 060/150] [Bugfix] Fix Ray compiled-DAG SHM channel stalls by detaching zero-copy `np.ndarray` logprobs buffers (#35736) Signed-off-by: Nikita Shapovalov --- tests/v1/executor/test_ray_utils.py | 54 +++++++++++++++++++++++++++++ vllm/v1/executor/ray_executor.py | 10 ++++-- vllm/v1/executor/ray_utils.py | 47 +++++++++++++++++++++++++ 3 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 tests/v1/executor/test_ray_utils.py diff --git a/tests/v1/executor/test_ray_utils.py b/tests/v1/executor/test_ray_utils.py new file mode 100644 index 00000000000..8da9d5459e7 --- /dev/null +++ b/tests/v1/executor/test_ray_utils.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np + +from vllm.v1.executor.ray_utils import detach_zero_copy_from_model_runner_output +from vllm.v1.outputs import LogprobsLists, LogprobsTensors, ModelRunnerOutput + + +def _make_readonly(arr: np.ndarray) -> np.ndarray: + arr.setflags(write=False) + return arr + + +def test_detach_zero_copy_from_model_runner_output_copies_only_numpy_views(): + cu_num_generated_tokens = [0, 2] + prompt_logprobs = LogprobsTensors.empty_cpu(1, 2) + output = ModelRunnerOutput( + req_ids=["req-0"], + req_id_to_index={"req-0": 0}, + logprobs=LogprobsLists( + logprob_token_ids=_make_readonly( + np.array([[1, 2], [3, 4]], dtype=np.int32) + ), + logprobs=_make_readonly( + np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32) + ), + sampled_token_ranks=_make_readonly(np.array([1, 2], dtype=np.int32)), + cu_num_generated_tokens=cu_num_generated_tokens, + ), + prompt_logprobs_dict={"req-0": prompt_logprobs}, + ) + + original_logprobs = output.logprobs + assert original_logprobs is not None + + detach_zero_copy_from_model_runner_output(output) + + detached_logprobs = output.logprobs + assert detached_logprobs is not None + assert detached_logprobs is not original_logprobs + assert ( + detached_logprobs.logprob_token_ids is not original_logprobs.logprob_token_ids + ) + assert detached_logprobs.logprobs is not original_logprobs.logprobs + assert ( + detached_logprobs.sampled_token_ranks + is not original_logprobs.sampled_token_ranks + ) + assert detached_logprobs.logprob_token_ids.flags.writeable + assert detached_logprobs.logprobs.flags.writeable + assert detached_logprobs.sampled_token_ranks.flags.writeable + assert detached_logprobs.cu_num_generated_tokens is cu_num_generated_tokens + assert output.prompt_logprobs_dict["req-0"] is prompt_logprobs diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py index 1dda2c294e4..cfeebb5e09d 100644 --- a/vllm/v1/executor/ray_executor.py +++ b/vllm/v1/executor/ray_executor.py @@ -26,6 +26,7 @@ from vllm.v1.executor.ray_utils import ( WORKER_SPECIFIC_ENV_VARS, FutureWrapper, RayWorkerWrapper, + detach_zero_copy_from_model_runner_output, initialize_ray_cluster, ray, ) @@ -463,7 +464,9 @@ class RayDistributedExecutor(Executor): # Get output only from a single worker (output_rank) # When PP is not used, we block here until the result is available. if not non_block: - return refs[0].get() + output = refs[0].get() + detach_zero_copy_from_model_runner_output(output) + return output # When PP is used, we return a FutureWrapper immediately so that # the scheduler can yield to the next batch. @@ -473,7 +476,10 @@ class RayDistributedExecutor(Executor): assert self.kv_output_aggregator is not None if not non_block: # Block and get results from all workers - return self.kv_output_aggregator.aggregate(ray.get(refs)) + outputs = ray.get(refs) + for output in outputs: + detach_zero_copy_from_model_runner_output(output) + return self.kv_output_aggregator.aggregate(outputs) # Return a future that will aggregate outputs from all workers return FutureWrapper(refs, self.kv_output_aggregator) diff --git a/vllm/v1/executor/ray_utils.py b/vllm/v1/executor/ray_utils.py index 67b43bdc11d..1541b24deaa 100644 --- a/vllm/v1/executor/ray_utils.py +++ b/vllm/v1/executor/ray_utils.py @@ -7,6 +7,8 @@ from collections import defaultdict from concurrent.futures import Future from typing import TYPE_CHECKING, Union +import numpy as np + import vllm.platforms from vllm.config import ParallelConfig from vllm.distributed import get_pp_group @@ -190,6 +192,48 @@ except ImportError as e: RayWorkerWrapper = None # type: ignore +def detach_zero_copy_from_model_runner_output(output: "ModelRunnerOutput") -> None: + """Detach Ray SHM-channel zero-copy buffers from a ModelRunnerOutput in-place. + + Ray compiled DAG SHM channels may return zero-copy objects (e.g. `np.ndarray`) + backed by Ray's shared-memory object store. Ray's channel docs explicitly + warn that subsequent reads may block if such an object is still in scope. + + vLLM can return numpy-backed logprobs in `ModelRunnerOutput.logprobs`. If + those arrays are backed by Ray SHM (commonly read-only), retaining them in + scope across scheduler iterations can stall the channel and eventually hit + `RAY_CGRAPH_get_timeout`. + + Copy read-only numpy arrays so the returned output no longer retains + references to Ray's shared-memory buffers. + + We intentionally do not touch `prompt_logprobs_dict`: those entries are + `LogprobsTensors` backed by PyTorch-owned CPU tensors (`to_cpu_nonblocking` + or `empty_cpu`), not NumPy views decoded from Ray channels. + """ + if output.logprobs is None: + return + + token_ids, logprobs, ranks, cu_num_generated_tokens = output.logprobs + + def _copy_if_readonly(arr): + if isinstance(arr, np.ndarray) and not arr.flags.writeable: + return arr.copy() + return arr + + # `cu_num_generated_tokens` is already a plain Python list (or None), so it + # never aliases Ray SHM buffers and can be reused as-is. + token_ids_c = _copy_if_readonly(token_ids) + logprobs_c = _copy_if_readonly(logprobs) + ranks_c = _copy_if_readonly(ranks) + if token_ids_c is token_ids and logprobs_c is logprobs and ranks_c is ranks: + return + + output.logprobs = type(output.logprobs)( + token_ids_c, logprobs_c, ranks_c, cu_num_generated_tokens + ) + + class FutureWrapper(Future): """A wrapper around Ray output reference to meet the interface of .execute_model(): The top level (core busy loop) expects .result() api @@ -207,8 +251,11 @@ class FutureWrapper(Future): def result(self, timeout=None): outputs = ray.get(self.ref_or_refs, timeout=timeout) if self.aggregator is None: + detach_zero_copy_from_model_runner_output(outputs) return outputs + for output in outputs: + detach_zero_copy_from_model_runner_output(output) return self.aggregator.aggregate(outputs, output_rank=0) From 617d1c2ff14d3b8b8d33ccb51f3235c3665adecb Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Thu, 16 Apr 2026 23:52:37 +0800 Subject: [PATCH 061/150] [Misc] Move `pyav` and `soundfile` to common requirements (#39997) Signed-off-by: Isotr0py --- requirements/common.txt | 2 ++ requirements/test/rocm.txt | 5 ++++- setup.py | 2 -- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/requirements/common.txt b/requirements/common.txt index 6e7fd90d023..6a183e09acf 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -32,7 +32,9 @@ pyzmq >= 25.0.0 msgspec gguf >= 0.17.0 mistral_common[image] >= 1.11.0 +av # required for audio in video IO opencv-python-headless >= 4.13.0 # required for video IO +soundfile # required for audio IO pyyaml six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12 setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12 diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 2e03df1bdc0..eaa11451a21 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -76,7 +76,9 @@ attrs==26.1.0 audioread==3.0.1 # via librosa av==16.1.0 - # via -r requirements/test/rocm.in + # via + # -r requirements/test/../common.txt + # -r requirements/test/rocm.in azure-core==1.39.0 # via # azure-identity @@ -1331,6 +1333,7 @@ sortedcontainers==2.4.0 # via hypothesis soundfile==0.13.1 # via + # -r requirements/test/../common.txt # -r requirements/test/rocm.in # genai-perf # librosa diff --git a/setup.py b/setup.py index 37782062ae0..7665978c6b0 100644 --- a/setup.py +++ b/setup.py @@ -1092,9 +1092,7 @@ setup( "instanttensor": ["instanttensor >= 0.1.5"], "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"], "audio": [ - "av", "scipy", - "soundfile", "mistral_common[audio]", ], # Required for audio processing "video": [], # Kept for backwards compatibility From 3abb7560c0d6f8e8e84a5c65eca68b9aa0c71dbb Mon Sep 17 00:00:00 2001 From: Roger Wang Date: Thu, 16 Apr 2026 11:53:58 -0700 Subject: [PATCH 062/150] [Bugfix] Fix audioflamingo test (#40052) Signed-off-by: Roger Wang --- .../processing/test_audioflamingo3.py | 85 ------------------- 1 file changed, 85 deletions(-) diff --git a/tests/models/multimodal/processing/test_audioflamingo3.py b/tests/models/multimodal/processing/test_audioflamingo3.py index 24311e5212b..3a6da75ed5b 100644 --- a/tests/models/multimodal/processing/test_audioflamingo3.py +++ b/tests/models/multimodal/processing/test_audioflamingo3.py @@ -140,88 +140,3 @@ def test_audio_token_count_matches_hf_processor_math(): _count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 0) == 1499 ) assert _count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 1) == 375 - - -def test_audio_feature_pipeline_matches_hf_small_config(): - from transformers.models.audioflamingo3 import ( - modeling_audioflamingo3 as hf_audioflamingo3_modeling, - ) - from transformers.models.audioflamingo3.configuration_audioflamingo3 import ( - AudioFlamingo3Config, - ) - - from vllm.model_executor.models.audioflamingo3 import ( - AudioFlamingo3Encoder, - AudioFlamingo3MultiModalProjector, - _build_audio_encoder_attention_mask, - _flatten_valid_audio_embeddings, - ) - - text_config = { - "model_type": "qwen2", - "intermediate_size": 64, - "initializer_range": 0.02, - "hidden_size": 32, - "max_position_embeddings": 1024, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "vocab_size": 128, - "pad_token_id": 1, - "use_mrope": False, - } - audio_config = { - "hidden_size": 16, - "num_attention_heads": 4, - "intermediate_size": 32, - "num_hidden_layers": 2, - "num_mel_bins": 80, - "max_source_positions": 1500, - "dropout": 0.0, - "attention_dropout": 0.0, - "activation_dropout": 0.0, - "encoder_layerdrop": 0.0, - } - - torch.manual_seed(0) - config = AudioFlamingo3Config( - text_config=text_config, - audio_config=audio_config, - audio_token_id=0, - ) - hf_model = hf_audioflamingo3_modeling.AudioFlamingo3ForConditionalGeneration( - config - ).eval() - - vllm_encoder = AudioFlamingo3Encoder(config.audio_config).eval() - vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict()) - - vllm_projector = AudioFlamingo3MultiModalProjector(config).eval() - vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict()) - - input_features = torch.randn(3, 80, 3000) - feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool) - feature_attention_mask[0, :3000] = True - feature_attention_mask[1, :2500] = True - feature_attention_mask[2, :1500] = True - - hf_output = hf_model.get_audio_features( - input_features, - feature_attention_mask, - return_dict=True, - ).pooler_output - vllm_attention_mask = _build_audio_encoder_attention_mask( - feature_attention_mask, - dtype=vllm_encoder.conv1.weight.dtype, - device=vllm_encoder.conv1.weight.device, - ) - vllm_hidden_states = vllm_encoder( - input_features, - attention_mask=vllm_attention_mask, - ) - vllm_output, _ = _flatten_valid_audio_embeddings( - vllm_projector(vllm_hidden_states), - feature_attention_mask, - ) - - torch.testing.assert_close(vllm_output, hf_output) From afabb5f45a6ffa3d9cb2f7424db984a9b8290098 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Fri, 17 Apr 2026 02:54:39 +0800 Subject: [PATCH 063/150] [bugfix] Normalize tool message content from array to string format (#39899) Signed-off-by: JaredforReal --- vllm/entrypoints/chat_utils.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index 3710473560d..c7c8c042169 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -1550,6 +1550,18 @@ def _parse_chat_message_content( parsed_msg = _ToolParser(message) if "tool_call_id" in parsed_msg: result_msg["tool_call_id"] = parsed_msg["tool_call_id"] + # Normalize tool message content from OpenAI array format to plain + # string. Clients like Claude Code / Cursor send tool results as + # [{"type": "text", "text": "..."}], but most chat templates only + # handle string content for tool messages. + msg_content = result_msg.get("content") + if isinstance(msg_content, list): + texts = [ + item.get("text", "") + for item in msg_content + if isinstance(item, dict) and item.get("type") == "text" + ] + result_msg["content"] = "\n".join(texts) if texts else "" if "name" in message and isinstance(message["name"], str): result_msg["name"] = message["name"] From de111f32461bf751851c175af59e913fafaa16d1 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Fri, 17 Apr 2026 03:01:25 +0800 Subject: [PATCH 064/150] [Bugfix] Fix bench_serve UTF-8 decode crash on split multi-byte chars (#38732) --- vllm/benchmarks/lib/endpoint_request_func.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index 61af098f80d..9c85ad686d3 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """The request function for API endpoints.""" +import codecs import io import json import os @@ -25,11 +26,12 @@ class StreamedResponseHandler: def __init__(self): self.buffer = "" + self._decoder = codecs.getincrementaldecoder("utf-8")() def add_chunk(self, chunk_bytes: bytes) -> list[str]: """Add a chunk of bytes to the buffer and return any complete messages.""" - chunk_str = chunk_bytes.decode("utf-8") + chunk_str = self._decoder.decode(chunk_bytes) self.buffer += chunk_str messages = [] From b16fda62b70af82d11f6637810f40f79ff713a82 Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Thu, 16 Apr 2026 15:25:29 -0400 Subject: [PATCH 065/150] [Misc] Add @sfeng33 to CODEOWNERS (#40048) Signed-off-by: sfeng33 <4florafeng@gmail.com> --- .github/CODEOWNERS | 10 +++++++--- docs/governance/committers.md | 3 ++- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index e1a2a582083..e272bf9c477 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -44,8 +44,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson /vllm/pooling_params.py @noooop @DarkLight1337 /vllm/tokenizers @DarkLight1337 @njhill /vllm/renderers @DarkLight1337 @njhill -/vllm/reasoning @aarnphm @chaunceyjiang -/vllm/tool_parsers @aarnphm @chaunceyjiang +/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 +/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 +/vllm/parser @aarnphm @chaunceyjiang @sfeng33 # vLLM V1 /vllm/v1/attention @LucasWilkinson @MatthewBonanni @@ -91,7 +92,10 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson /tests/v1/kv_connector/nixl_integration @NickLucche /tests/v1/kv_connector @ApostaC @orozery /tests/v1/kv_offload @ApostaC @orozery -/tests/v1/determinism @yewentao256 +/tests/v1/determinism @yewentao256 +/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 +/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 +/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 # Transformers modeling backend /vllm/model_executor/models/transformers @hmellor diff --git a/docs/governance/committers.md b/docs/governance/committers.md index e6f8f317ebb..fb0d09a2a9d 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -41,6 +41,7 @@ Sorted alphabetically by GitHub handle: - [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg - [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support - [@russellb](https://github.com/russellb): Structured output, engine core, security +- [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser - [@sighingnow](https://github.com/sighingnow): Qwen models, new model support - [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community - [@tdoublep](https://github.com/tdoublep): State space models @@ -119,7 +120,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - State space models: The state space models implementation in vLLM - @tdoublep, @tlrmchlsmth - Reasoning and tool calling parsers - - @chaunceyjiang, @aarnphm + - @chaunceyjiang, @aarnphm, @sfeng33 ### Entrypoints From adf9bb3c577aaaad147b1fe7f61a5c6a0bdfb3de Mon Sep 17 00:00:00 2001 From: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com> Date: Thu, 16 Apr 2026 12:51:45 -0700 Subject: [PATCH 066/150] [CI] Add weight transfer tests to CI (#39821) Signed-off-by: SumanthRH Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Cyrus Leung --- .buildkite/test_areas/distributed.yaml | 2 ++ tests/distributed/test_weight_transfer.py | 3 +++ 2 files changed, 5 insertions(+) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 56e528ffb37..e13618eb65d 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -196,6 +196,8 @@ steps: - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - pytest -v -s tests/v1/distributed/test_dbo.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py + - pytest -v -s tests/distributed/test_packed_tensor.py - label: Distributed Tests (2 GPUs)(B200) device: b200 diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 1c9bc766ab1..92b8100a166 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -41,6 +41,7 @@ def create_mock_parallel_config( config.rank = rank config.world_size = world_size config.data_parallel_rank = dp_rank + config.data_parallel_index = dp_rank return config @@ -283,6 +284,7 @@ def inference_receive_tensor( parallel_config.rank = 0 parallel_config.world_size = 1 parallel_config.data_parallel_rank = 0 + parallel_config.data_parallel_index = 0 engine = NCCLWeightTransferEngine(config, parallel_config) @@ -666,6 +668,7 @@ def inference_receive_ipc_tensor( parallel_config.rank = 0 parallel_config.world_size = 1 parallel_config.data_parallel_rank = 0 + parallel_config.data_parallel_index = 0 engine = IPCWeightTransferEngine(config, parallel_config) From b897f00c9c35dcc7b229973cf665a49d7082b8bf Mon Sep 17 00:00:00 2001 From: roikoren755 <26850796+roikoren755@users.noreply.github.com> Date: Thu, 16 Apr 2026 23:06:01 +0300 Subject: [PATCH 067/150] Gate SSU dispatch setup (#40039) Signed-off-by: Roi Koren --- tests/kernels/mamba/test_ssu_dispatch.py | 52 +++++++++++++++++-- .../layers/mamba/ops/ssu_dispatch.py | 23 ++++++-- vllm/v1/worker/gpu/model_runner.py | 4 +- vllm/v1/worker/gpu_model_runner.py | 4 +- 4 files changed, 73 insertions(+), 10 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 96b04f44d22..887d60b2736 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -13,6 +13,11 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( selective_state_update, ) from vllm.utils.torch_utils import set_random_seed +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) try: import flashinfer.mamba # noqa: F401 @@ -22,22 +27,40 @@ except ImportError: HAS_FLASHINFER = False +def _kv_cache_config_with_ssu(mamba_type: str = "mamba2") -> KVCacheConfig: + spec = MambaSpec( + block_size=16, + shapes=((16, 64),), + dtypes=(torch.float16,), + mamba_type=mamba_type, + ) + return KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["l0"], kv_cache_spec=spec)], + ) + + def test_default_backend_is_triton(): - initialize_mamba_ssu_backend(MambaConfig()) + initialize_mamba_ssu_backend(MambaConfig(), _kv_cache_config_with_ssu()) backend = get_mamba_ssu_backend() assert isinstance(backend, TritonSSUBackend) assert backend.name == "triton" def test_explicit_triton_backend(): - initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON)) + initialize_mamba_ssu_backend( + MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu() + ) backend = get_mamba_ssu_backend() assert isinstance(backend, TritonSSUBackend) @pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed") def test_flashinfer_backend_init(): - initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER)) + initialize_mamba_ssu_backend( + MambaConfig(backend=MambaBackendEnum.FLASHINFER), _kv_cache_config_with_ssu() + ) backend = get_mamba_ssu_backend() assert isinstance(backend, FlashInferSSUBackend) assert backend.name == "flashinfer" @@ -53,6 +76,25 @@ def test_uninitialized_backend_raises(): mod._mamba_ssu_backend = old +@pytest.mark.parametrize( + "mamba_type", ["linear_attention", "gdn_attention", "short_conv"] +) +def test_init_is_noop_for_non_ssu_mamba_type(mamba_type): + import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod + + old = mod._mamba_ssu_backend + mod._mamba_ssu_backend = None + try: + initialize_mamba_ssu_backend( + MambaConfig(), _kv_cache_config_with_ssu(mamba_type) + ) + assert mod._mamba_ssu_backend is None + with pytest.raises(RuntimeError, match="not been initialized"): + get_mamba_ssu_backend() + finally: + mod._mamba_ssu_backend = old + + @pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed") def test_flashinfer_import_error(): with pytest.raises(ImportError, match="FlashInfer is required"): @@ -61,7 +103,9 @@ def test_flashinfer_import_error(): def test_triton_basic_call(): set_random_seed(0) - initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON)) + initialize_mamba_ssu_backend( + MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu() + ) device = "cuda" batch_size = 2 dim = 64 diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 8a86b1a068b..33a08feb9cf 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -15,6 +15,7 @@ import torch from vllm.config.mamba import MambaBackendEnum, MambaConfig from vllm.logger import init_logger from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec logger = init_logger(__name__) @@ -188,12 +189,22 @@ _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = { _mamba_ssu_backend: MambaSSUBackend | None = None -def initialize_mamba_ssu_backend(mamba_config: MambaConfig) -> None: +def initialize_mamba_ssu_backend( + mamba_config: MambaConfig, + kv_cache_config: KVCacheConfig, +) -> None: """Initialize the global Mamba SSU backend. - Args: - mamba_config: Mamba configuration. + No-op if `kv_cache_config` contains no specs that call + selective_state_update. """ + if not any( + isinstance(g.kv_cache_spec, MambaSpec) + and g.kv_cache_spec.mamba_type in ("mamba1", "mamba2") + for g in kv_cache_config.kv_cache_groups + ): + return + global _mamba_ssu_backend backend = mamba_config.backend @@ -203,7 +214,11 @@ def initialize_mamba_ssu_backend(mamba_config: MambaConfig) -> None: f"Valid options: {list(_BACKEND_REGISTRY.keys())}" ) - _mamba_ssu_backend = _BACKEND_REGISTRY[backend](mamba_config) + backend_cls = _BACKEND_REGISTRY[backend] + if isinstance(_mamba_ssu_backend, backend_cls): + return + + _mamba_ssu_backend = backend_cls(mamba_config) logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 9b79a13ee38..a0025d8c795 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -363,7 +363,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.attn_backends, self.attn_groups, attn_cg_support = init_attn_backend( self.kv_cache_config, self.vllm_config, self.device ) - initialize_mamba_ssu_backend(self.vllm_config.mamba_config) + initialize_mamba_ssu_backend( + self.vllm_config.mamba_config, self.kv_cache_config + ) cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( attn_cg_support.min_cg_support, attn_cg_support.min_cg_attn_backend, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 4c573f79e97..d95e70d071a 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6738,7 +6738,9 @@ class GPUModelRunner( self.may_add_encoder_only_layers_to_kv_cache_config() self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config) self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling) - initialize_mamba_ssu_backend(self.vllm_config.mamba_config) + initialize_mamba_ssu_backend( + self.vllm_config.mamba_config, self.kv_cache_config + ) # The kernel block size for all KV cache groups. For example, if # kv_cache_manager uses block_size 256 for a given group, but the attention # backends for that group only supports block_size 64, we will return From ad2b1277f99f1cee3d59d687aa1813267299f6a2 Mon Sep 17 00:00:00 2001 From: Asaf Gardin <39553475+Josephasafg@users.noreply.github.com> Date: Thu, 16 Apr 2026 23:12:20 +0300 Subject: [PATCH 068/150] [Quantization] Consolidate experts_int8 with fp8 online quantization (#38463) Signed-off-by: Josephasafg --- tests/quantization/test_experts_int8.py | 1 - vllm/config/quantization.py | 4 + .../layers/fused_moe/oracle/int8.py | 84 +++++++++ .../layers/quantization/__init__.py | 2 +- .../layers/quantization/experts_int8.py | 168 ++--------------- .../layers/quantization/online/base.py | 18 +- .../layers/quantization/online/fp8.py | 158 +--------------- .../layers/quantization/online/int8.py | 109 +++++++++++ .../layers/quantization/online/moe_base.py | 172 ++++++++++++++++++ 9 files changed, 403 insertions(+), 313 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/oracle/int8.py create mode 100644 vllm/model_executor/layers/quantization/online/int8.py create mode 100644 vllm/model_executor/layers/quantization/online/moe_base.py diff --git a/tests/quantization/test_experts_int8.py b/tests/quantization/test_experts_int8.py index 22edb9c58da..7cdb135fa07 100644 --- a/tests/quantization/test_experts_int8.py +++ b/tests/quantization/test_experts_int8.py @@ -38,6 +38,5 @@ def test_model_experts_int8_startup( dtype=dtype, enforce_eager=True, quantization="experts_int8", - allow_deprecated_quantization=True, ) as vllm_model: vllm_model.generate_greedy(example_prompts, max_tokens) diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index 1b7022380ee..02df08b8b48 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -19,6 +19,10 @@ class OnlineQuantScheme(Enum): # blocks of 128x128 elements (popularized by DeepSeek) FP8_PER_BLOCK = "fp8_per_block" + # int8, weight-only per-channel quantization for MoE expert weights. + # Linear layers remain unquantized. + INT8_PER_CHANNEL_WEIGHT_ONLY = "int8_per_channel_weight_only" + # TODO(future PRs): add more online quant schemes here: mxfp8, etc diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py new file mode 100644 index 00000000000..3ae9a491e9b --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, + int8_w8a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) + +logger = init_logger(__name__) + + +def select_int8_moe_backend( + config: FusedMoEConfig, +) -> type[mk.FusedMoEExperts]: + from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts + + supported, reason = TritonExperts.is_supported_config( + TritonExperts, + config, + None, + None, + mk.FusedMoEActivationFormat.Standard, + ) + if not supported: + raise ValueError( + f"INT8 Triton MoE backend does not support the " + f"deployment configuration: {reason}" + ) + + logger.info_once("Using Triton INT8 MoE backend", scope="local") + return TritonExperts + + +def make_int8_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, +) -> FusedMoEQuantConfig: + return int8_w8a16_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_zp=None, + w2_zp=None, + ) + + +def make_int8_moe_kernel( + moe_quant_config: FusedMoEQuantConfig, + moe_config: FusedMoEConfig, + experts_cls: type[mk.FusedMoEExperts], + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + shared_experts: SharedExperts | None = None, +) -> mk.FusedMoEKernel: + prepare_finalize = maybe_make_prepare_finalize( + moe=moe_config, + quant_config=moe_quant_config, + routing_tables=routing_tables, + allow_new_interface=True, + ) + assert prepare_finalize is not None + + logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local") + + experts = experts_cls( + moe_config=moe_config, + quant_config=moe_quant_config, + ) + + return mk.FusedMoEKernel( + prepare_finalize, + experts, + shared_experts=shared_experts, + inplace=not moe_config.disable_inplace, + ) diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index 6313db78a82..352d04ea208 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -40,6 +40,7 @@ QuantizationMethods = Literal[ # shorthand for creating a more complicated online quant config object "fp8_per_tensor", "fp8_per_block", + "int8_per_channel_weight_only", ] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) @@ -47,7 +48,6 @@ DEPRECATED_QUANTIZATION_METHODS = [ "tpu_int8", "fbgemm_fp8", "fp_quant", - "experts_int8", ] # The customized quantization methods which will be added to this dict. diff --git a/vllm/model_executor/layers/quantization/experts_int8.py b/vllm/model_executor/layers/quantization/experts_int8.py index 301441ff019..007c3c214ee 100644 --- a/vllm/model_executor/layers/quantization/experts_int8.py +++ b/vllm/model_executor/layers/quantization/experts_int8.py @@ -5,27 +5,25 @@ from typing import Any import torch -from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, - FusedMoEConfig, - FusedMoEMethodBase, -) -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEQuantConfig, - int8_w8a16_moe_quant_config, -) +from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod from vllm.model_executor.layers.quantization import QuantizationMethods from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, ) -from vllm.model_executor.utils import set_weight_attrs +from vllm.model_executor.layers.quantization.online.int8 import ( + Int8OnlineMoEMethod, +) class ExpertsInt8Config(QuantizationConfig): - """Config class for Int8 experts quantization.""" + """Online int8 quantization for MoE expert weights. + Linear layers are left unquantized. + + Backward-compatible config for ``--quantization experts_int8``. + Prefer ``--quantization int8_per_channel`` + """ def __init__(self) -> None: super().__init__() @@ -56,149 +54,5 @@ class ExpertsInt8Config(QuantizationConfig): if isinstance(layer, LinearBase): return UnquantizedLinearMethod() elif isinstance(layer, FusedMoE): - return ExpertsInt8MoEMethod(self, layer.moe_config) + return Int8OnlineMoEMethod(layer=layer) return None - - -class ExpertsInt8MoEMethod(FusedMoEMethodBase): - def __init__( - self, - quant_config: ExpertsInt8Config, - moe: FusedMoEConfig, - ): - super().__init__(moe) - self.quant_config = quant_config - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - int8_dtype = torch.int8 - - assert "weight_loader" in extra_weight_attrs - weight_loader = extra_weight_attrs["weight_loader"] - wrapped_weight_loader = ExpertsInt8MoEMethod.quantizing_weight_loader( - layer, weight_loader - ) - extra_weight_attrs["weight_loader"] = wrapped_weight_loader - - # Fused gate_up_proj (column parallel) - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size, - dtype=int8_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - - # down_proj (row parallel) - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition, - dtype=int8_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - - w13_scale = torch.nn.Parameter( - torch.zeros( - num_experts, 2 * intermediate_size_per_partition, dtype=torch.float32 - ), - requires_grad=False, - ) - layer.register_parameter("w13_scale", w13_scale) - - w2_scale = torch.nn.Parameter( - torch.zeros(num_experts, hidden_size, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_scale", w2_scale) - - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - return int8_w8a16_moe_quant_config( - w1_scale=layer.w13_scale, w2_scale=layer.w2_scale, w1_zp=None, w2_zp=None - ) - - def apply( - self, - layer: FusedMoE, - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - from vllm.model_executor.layers.fused_moe import fused_experts - - return fused_experts( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights=topk_weights, - topk_ids=topk_ids, - inplace=not self.moe.disable_inplace, - activation=layer.activation, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - quant_config=self.moe_quant_config, - ) - - @staticmethod - def quantizing_weight_loader(layer, weight_loader): - def quantize_and_call_weight_loader( - param: torch.nn.Parameter, - loaded_weight: torch.Tensor, - weight_name: str, - shard_id: int, - expert_id: int, - ): - tp_rank = get_tensor_model_parallel_rank() - shard_size = layer.intermediate_size_per_partition - shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size) - device = get_tp_group().device - loaded_weight = loaded_weight.to(device) - # w1, gate_proj case: Load into first shard of w13. - if shard_id == "w1": - scales = quantize_in_place_and_get_scales(loaded_weight[shard, :]) - layer.w13_scale.data[expert_id, 0:shard_size].copy_(scales[:, 0]) - # w3, up_proj case: Load into second shard of w13. - elif shard_id == "w3": - scales = quantize_in_place_and_get_scales(loaded_weight[shard, :]) - layer.w13_scale.data[expert_id, shard_size : 2 * shard_size].copy_( - scales[:, 0] - ) - # w2, down_proj case: Load into only shard of w2. - elif shard_id == "w2": - scales = quantize_in_place_and_get_scales(loaded_weight[:, shard]) - layer.w2_scale.data[expert_id, :].copy_(scales[:, 0]) - else: - raise ValueError(f"Shard id must be in [0,1,2] but got {shard_id}") - weight_loader(param, loaded_weight, weight_name, shard_id, expert_id) - - return quantize_and_call_weight_loader - - -def quantize_in_place_and_get_scales(weight: torch.Tensor) -> torch.Tensor: - vmax = torch.iinfo(torch.int8).max - scales = torch.max(torch.abs(weight), dim=1, keepdim=True)[0] / vmax - - weight.div_(scales) - weight.round_() - weight.clamp_(-vmax, vmax) - - return scales diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index 87997f8ef68..1f923c07ff1 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -9,6 +9,7 @@ from vllm.config.quantization import ( OnlineQuantizationConfigArgs, OnlineQuantScheme, ) +from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( FusedMoE, ) @@ -33,6 +34,11 @@ from vllm.model_executor.layers.quantization.online.fp8 import ( Fp8PerTensorOnlineLinearMethod, Fp8PerTensorOnlineMoEMethod, ) +from vllm.model_executor.layers.quantization.online.int8 import ( + Int8OnlineMoEMethod, +) + +logger = init_logger(__name__) class OnlineQuantizationConfig(QuantizationConfig): @@ -96,7 +102,13 @@ class OnlineQuantizationConfig(QuantizationConfig): return UnquantizedLinearMethod() linear_scheme = self.args.linear_scheme_override or self.args.global_scheme - if linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK: + if linear_scheme == OnlineQuantScheme.INT8_PER_CHANNEL_WEIGHT_ONLY: + logger.warning_once( + "INT8 online quantization only quantizes MoE expert " + "weights. linear layers remain in full precision." + ) + return UnquantizedLinearMethod() + elif linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK: return Fp8PerBlockOnlineLinearMethod() else: return Fp8PerTensorOnlineLinearMethod() @@ -109,7 +121,9 @@ class OnlineQuantizationConfig(QuantizationConfig): return UnquantizedFusedMoEMethod(layer.moe_config) moe_scheme = self.args.moe_scheme_override or self.args.global_scheme - if moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK: + if moe_scheme == OnlineQuantScheme.INT8_PER_CHANNEL_WEIGHT_ONLY: + return Int8OnlineMoEMethod(layer=layer) + elif moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK: return Fp8PerBlockOnlineMoEMethod(layer=layer) else: return Fp8PerTensorOnlineMoEMethod(layer=layer) diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index fa8cf240627..9cb697289d7 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -10,7 +10,6 @@ if TYPE_CHECKING: import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend @@ -19,15 +18,15 @@ import vllm.envs as envs from vllm import _custom_ops as ops from vllm.config import get_current_vllm_config from vllm.model_executor.kernels.linear import init_fp8_linear_kernel -from vllm.model_executor.layers.fused_moe import ( - FusedMoEMethodBase, -) from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( select_fp8_moe_backend, ) from vllm.model_executor.layers.linear import ( LinearMethodBase, ) +from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, create_fp8_quant_key, @@ -44,7 +43,7 @@ from vllm.model_executor.model_loader.reload.layerwise import ( initialize_online_processing, ) from vllm.model_executor.parameter import ModelWeightParameter -from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.utils.deep_gemm import per_block_cast_to_fp8 @@ -268,21 +267,15 @@ class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase): # --------------------------------------------------------------------------- -class _Fp8OnlineMoEBase(FusedMoEMethodBase): +class _Fp8OnlineMoEBase(OnlineMoEMethodBase): """Shared base for online FP8 MoE methods. Loads fp16/bf16 checkpoint weights onto meta device and materializes them just-in-time.""" - uses_meta_device: bool = True - # Declared here for mypy; actual values are set in __init__. fp8_backend: "Fp8MoeBackend" experts_cls: "type[mk.FusedMoEExperts] | None" weight_scale_name: str weight_block_size: list[int] | None - moe: "FusedMoEConfig" - is_monolithic: bool - moe_quant_config: "FusedMoEQuantConfig | None" - moe_kernel: "mk.FusedMoEKernel | None" def __init__( self, @@ -313,77 +306,6 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase): allow_vllm_cutlass=False, ) - def create_weights( - self, - layer: Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - layer.num_experts = num_experts - layer.orig_dtype = params_dtype - layer.weight_block_size = None - - # WEIGHTS - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size, - device="meta", - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition, - device="meta", # materialized and processed during loading - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - - # BIASES (for models like GPT-OSS that have biased MoE) - if self.moe.has_bias: - w13_bias = torch.nn.Parameter( - torch.zeros( - num_experts, - 2 * intermediate_size_per_partition, - device="meta", # materialized and processed during loading - dtype=layer.orig_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_bias", w13_bias) - set_weight_attrs(w13_bias, extra_weight_attrs) - - w2_bias = torch.nn.Parameter( - torch.zeros( - num_experts, - hidden_size, - device="meta", # materialized and processed during loading - dtype=layer.orig_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_bias", w2_bias) - set_weight_attrs(w2_bias, extra_weight_attrs) - - layer.w13_input_scale = None - layer.w2_input_scale = None - - initialize_online_processing(layer) - def _setup_kernel( self, layer: "FusedMoE", @@ -430,15 +352,6 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase): shared_experts=layer.shared_experts, ) - def maybe_make_prepare_finalize( - self, - routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> "mk.FusedMoEPrepareAndFinalizeModular | None": - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel " - "initialization logic. This function should not be called." - ) - def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> "FusedMoEQuantConfig": @@ -460,68 +373,9 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase): block_shape=self.weight_block_size, ) - # Inject biases into the quant config if the model has them - # (e.g. GPT-OSS biased MoE) - if quant_config is not None and self.moe.has_bias: - w13_bias = getattr(layer, "w13_bias", None) - w2_bias = getattr(layer, "w2_bias", None) - if w13_bias is not None: - quant_config._w1.bias = w13_bias - if w2_bias is not None: - quant_config._w2.bias = w2_bias - + self._maybe_inject_biases(quant_config, layer) return quant_config - @property - def supports_eplb(self) -> bool: - return True - - def apply_monolithic( - self, - layer: "FusedMoE", - x: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.is_monolithic - assert self.moe_kernel is not None - return self.moe_kernel.apply_monolithic( - x, - layer.w13_weight, - layer.w2_weight, - router_logits, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - num_expert_group=layer.num_expert_group, - topk_group=layer.topk_group, - e_score_correction_bias=layer.e_score_correction_bias, - routed_scaling_factor=layer.routed_scaling_factor, - ) - - def apply( - self, - layer: "FusedMoE", - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert not self.is_monolithic - assert self.moe_kernel is not None - return self.moe_kernel.apply( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - shared_experts_input=shared_experts_input, - ) - class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase): """Online tensorwise FP8 MoE quantization. diff --git a/vllm/model_executor/layers/quantization/online/int8.py b/vllm/model_executor/layers/quantization/online/int8.py new file mode 100644 index 00000000000..18cc6aa1860 --- /dev/null +++ b/vllm/model_executor/layers/quantization/online/int8.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING + +import torch +from torch.nn import Module + +if TYPE_CHECKING: + import vllm.model_executor.layers.fused_moe.modular_kernel as mk + from vllm.model_executor.layers.fused_moe import FusedMoE + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, + ) + +from vllm.model_executor.layers.fused_moe.oracle.int8 import ( + make_int8_moe_kernel, + make_int8_moe_quant_config, + select_int8_moe_backend, +) +from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, +) +from vllm.model_executor.utils import replace_parameter + + +class Int8OnlineMoEMethod(OnlineMoEMethodBase): + """Online per-channel INT8 MoE quantization. + Loads fp16/bf16 weights and quantizes them per-row to int8 during loading. + """ + + def __init__( + self, + *, + layer: torch.nn.Module, + ): + super().__init__(layer.moe_config) + self.experts_cls: type[mk.FusedMoEExperts] = select_int8_moe_backend( + config=self.moe, + ) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + self._quantize_weights(layer) + self._setup_kernel(layer) + + layer._already_called_process_weights_after_loading = True + + def _quantize_weights(self, layer: Module) -> None: + vmax = torch.iinfo(torch.int8).max + + w13 = torch.empty_like(layer.w13_weight, dtype=torch.int8) + w2 = torch.empty_like(layer.w2_weight, dtype=torch.int8) + w13_scale = torch.zeros( + layer.num_experts, + layer.w13_weight.shape[1], + device=w13.device, + dtype=torch.float32, + ) + w2_scale = torch.zeros( + layer.num_experts, + layer.w2_weight.shape[1], + device=w2.device, + dtype=torch.float32, + ) + + for expert in range(layer.local_num_experts): + # w13: per-row quantization over hidden_size dim + w = layer.w13_weight[expert, :, :] + scales = w.abs().amax(dim=1) / vmax + q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax) + w13[expert, :, :] = q.to(torch.int8) + w13_scale[expert, :] = scales + + # w2: per-row quantization over intermediate_size dim + w = layer.w2_weight[expert, :, :] + scales = w.abs().amax(dim=1) / vmax + q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax) + w2[expert, :, :] = q.to(torch.int8) + w2_scale[expert, :] = scales + + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_scale", w13_scale) + replace_parameter(layer, "w2_scale", w2_scale) + + def _setup_kernel(self, layer: "FusedMoE") -> None: + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + assert self.experts_cls is not None + self.moe_kernel = make_int8_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + routing_tables=layer._maybe_init_expert_routing_tables(), + shared_experts=layer.shared_experts, + ) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> "FusedMoEQuantConfig | None": + quant_config = make_int8_moe_quant_config( + w1_scale=layer.w13_scale, + w2_scale=layer.w2_scale, + ) + self._maybe_inject_biases(quant_config, layer) + return quant_config diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py new file mode 100644 index 00000000000..25c3359ee8b --- /dev/null +++ b/vllm/model_executor/layers/quantization/online/moe_base.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import abstractmethod + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe import FusedMoEMethodBase +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.model_loader.reload.layerwise import ( + initialize_online_processing, +) +from vllm.model_executor.utils import set_weight_attrs + + +class OnlineMoEMethodBase(FusedMoEMethodBase): + """Base for MoE methods that load full-precision weights on meta device + and quantize them after loading via the QeRL layerwise processing system. + """ + + uses_meta_device: bool = True + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + layer.num_experts = num_experts + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + # Fused gate_up_proj (column parallel) — full precision on meta device + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + device="meta", + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + # down_proj (row parallel) — full precision on meta device + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + device="meta", + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # BIASES (for models like GPT-OSS that have biased MoE) + if self.moe.has_bias: + w13_bias = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + device="meta", + dtype=layer.orig_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_bias", w13_bias) + set_weight_attrs(w13_bias, extra_weight_attrs) + + w2_bias = torch.nn.Parameter( + torch.zeros( + num_experts, + hidden_size, + device="meta", + dtype=layer.orig_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_bias", w2_bias) + set_weight_attrs(w2_bias, extra_weight_attrs) + + layer.w13_input_scale = None + layer.w2_input_scale = None + + initialize_online_processing(layer) + + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + pass + + def _maybe_inject_biases( + self, + quant_config: FusedMoEQuantConfig, + layer: torch.nn.Module, + ) -> None: + """Inject biases into the quant config if the model has them + (e.g. GPT-OSS biased MoE).""" + if self.moe.has_bias: + w13_bias = getattr(layer, "w13_bias", None) + w2_bias = getattr(layer, "w2_bias", None) + if w13_bias is not None: + quant_config._w1.bias = w13_bias + if w2_bias is not None: + quant_config._w2.bias = w2_bias + + def maybe_make_prepare_finalize( + self, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: + raise ValueError( + f"{self.__class__.__name__} uses the new modular kernel " + "initialization logic. This function should not be called." + ) + + @property + def supports_eplb(self) -> bool: + return True + + def apply_monolithic( + self, + layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821 + x: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + ) + + def apply( + self, + layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821 + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts_input=shared_experts_input, + ) From 219bb5b8c0dcc6a5d5f894e9168fa5b8c2f8255a Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 16 Apr 2026 16:48:41 -0400 Subject: [PATCH 069/150] [Misc] Update `committers.md` (#40058) Signed-off-by: Matthew Bonanni --- docs/governance/committers.md | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/docs/governance/committers.md b/docs/governance/committers.md index fb0d09a2a9d..42c8a8ca2fe 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -31,6 +31,7 @@ Sorted alphabetically by GitHub handle: - [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance - [@luccafong](https://github.com/luccafong): Llama models, speculative decoding, distributed - [@markmc](https://github.com/markmc): Observability +- [@MatthewBonanni](https://github.com/MatthewBonanni): Kernels and performance - [@mgoin](https://github.com/mgoin): Quantization and performance - [@NickLucche](https://github.com/NickLucche): KV connector - [@njhill](https://github.com/njhill): Distributed, API server, engine core @@ -87,7 +88,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints - @robertgshaw2-redhat, @njhill, @russellb - ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation - - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg + - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni - KV Connector: Connector interface and implementation for KV cache offload and transfer - @robertgshaw2-redhat, @njhill, @KuntaiDu, @NickLucche, @ApostaC - Distributed, Parallelism, Process Management: Process launchers managing each worker, and assign them to the right DP/TP/PP/EP ranks @@ -106,7 +107,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - Custom Layers: Utility layers in vLLM such as rotary embedding and rms norms - @ProExpertProg - Attention: Attention interface for paged attention - - @WoosukKwon, @LucasWilkinson, @heheda12345 + - @WoosukKwon, @LucasWilkinson, @heheda12345, @MatthewBonanni - FusedMoE: FusedMoE kernel, Modular kernel framework, EPLB - @tlrmchlsmth - Quantization: Various quantization config, weight loading, and kernel. @@ -134,7 +135,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r ### Features - Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP. - - @WoosukKwon, @benchislett, @luccafong + - @WoosukKwon, @benchislett, @luccafong, @MatthewBonanni - Structured Output: The structured output implementation - @russellb, @aarnphm - RL: The RL related features such as collective rpc, sleep mode, etc. @@ -154,8 +155,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r ### External Kernels Integration -- FlashAttention: @LucasWilkinson -- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon +- FlashAttention: @LucasWilkinson, @MatthewBonanni +- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon, @MatthewBonanni - Blackwell Kernels: @mgoin, @yewentao256 - DeepEP/DeepGEMM: @mgoin, @yewentao256 From 29057d3bee1e4a5f84a41eb0cbd2f67b9fa35816 Mon Sep 17 00:00:00 2001 From: BadrBasowid <61441185+BadrBasowid@users.noreply.github.com> Date: Fri, 17 Apr 2026 06:57:16 +0800 Subject: [PATCH 070/150] [Compilation] Add Unit Tests for VllmFusionPatternMatcherPass (#39692) Signed-off-by: BadrBasowid --- .../test_vllm_fusion_pattern_matcher_pass.py | 128 ++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py diff --git a/tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py b/tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py new file mode 100644 index 00000000000..ee2eef009fc --- /dev/null +++ b/tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm.config +from tests.compile.backend import TestBackend +from vllm.platforms import current_platform +from vllm.compilation.passes.vllm_inductor_pass import ( + VllmFusionPatternMatcherPass, + VllmPatternMatcherPass, + VllmPatternReplacement, +) +from vllm.config import CompilationConfig, CompilationMode, VllmConfig + + + +class ReluToAbsPattern(VllmPatternReplacement): + """Replaces relu(x) with abs(x) — a minimal test fixture.""" + + @property + def pattern(self): + def _pattern(x: torch.Tensor) -> torch.Tensor: + return torch.ops.aten.relu.default(x) + + return _pattern + + @property + def replacement(self): + def _replacement(x: torch.Tensor) -> torch.Tensor: + return torch.ops.aten.abs.default(x) + + return _replacement + + def get_inputs(self) -> list[torch.Tensor]: + return [self.empty_fp32(4)] + + +class ExpToSqrtPattern(VllmPatternReplacement): + """A second distinct pattern type — used to test uuid differentiation.""" + + @property + def pattern(self): + def _pattern(x: torch.Tensor) -> torch.Tensor: + return torch.ops.aten.exp.default(x) + + return _pattern + + @property + def replacement(self): + def _replacement(x: torch.Tensor) -> torch.Tensor: + return torch.ops.aten.sqrt.default(x) + + return _replacement + + def get_inputs(self) -> list[torch.Tensor]: + return [self.empty_fp32(4)] + + + +class ReluFusionPass(VllmFusionPatternMatcherPass): + def __init__(self, config: VllmConfig) -> None: + super().__init__(config, "test_relu_fusion") + self.register(ReluToAbsPattern()) + + +class TwoPatternFusionPass(VllmFusionPatternMatcherPass): + def __init__(self, config: VllmConfig) -> None: + super().__init__(config, "test_two_pattern_fusion") + self.register(ReluToAbsPattern()) + self.register(ExpToSqrtPattern()) + + + +@pytest.fixture +def vllm_config(): + return VllmConfig( + compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE), + ) + +@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA") +def test_register_tracks_patterns(vllm_config): + """register() appends each VllmPatternReplacement to _pattern_replacements.""" + with vllm.config.set_current_vllm_config(vllm_config): + single = ReluFusionPass(vllm_config) + two = TwoPatternFusionPass(vllm_config) + + assert len(single._pattern_replacements) == 1 + assert len(two._pattern_replacements) == 2 + + +@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA") +def test_uuid_stable(vllm_config): + """Two instances of the same pass class produce identical uuids.""" + with vllm.config.set_current_vllm_config(vllm_config): + p1 = ReluFusionPass(vllm_config) + p2 = ReluFusionPass(vllm_config) + p3= TwoPatternFusionPass(vllm_config) + + assert p1.uuid() == p2.uuid() + assert p1.uuid() != p3.uuid() + assert p2.uuid() != p3.uuid() + + +@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA") +@pytest.mark.parametrize("N", [1, 2, 4]) +def test_matched_count_and_match_table(vllm_config, N): + """matched_count and match_table reflect the number of matched patterns.""" + + class Model(torch.nn.Module): + def forward(self, *inputs): + # N independent relus + return sum(torch.relu(x) for x in inputs) + + with vllm.config.set_current_vllm_config(vllm_config): + torch.set_default_device("cuda") + torch.set_default_dtype(torch.float32) + + fusion_pass = ReluFusionPass(vllm_config) + backend = TestBackend(fusion_pass) + model = torch.compile(Model(), backend=backend) + + inputs = [torch.rand(8) for _ in range(N)] + model(*inputs) + + assert fusion_pass.matched_count == N + assert VllmPatternMatcherPass.match_table["test_relu_fusion"] >= N From c4e601c73c2a2652ff7200d9f120efb7bde20faa Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Fri, 17 Apr 2026 02:22:05 +0300 Subject: [PATCH 071/150] Bugfix: Parakeet: `.conv.pointwise/depthwise_conv1/2.bias weigths` can exist even if `convolution_bias=False` (#40007) Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> --- vllm/model_executor/models/parakeet.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/vllm/model_executor/models/parakeet.py b/vllm/model_executor/models/parakeet.py index 40539cafb31..5671ba05c56 100644 --- a/vllm/model_executor/models/parakeet.py +++ b/vllm/model_executor/models/parakeet.py @@ -99,6 +99,8 @@ class ProjectedParakeet(nn.Module): if target is None: target = buffers_dict.get(target_name) if target is None: + if self._can_skip_missing_named_param(target_name): + continue raise ValueError(f"Unknown weight: {name}") weight_loader = getattr(target, "weight_loader", default_weight_loader) with torch.no_grad(): @@ -107,6 +109,27 @@ class ProjectedParakeet(nn.Module): return loaded_params + def _can_skip_missing_named_param(self, target_name: str) -> bool: + if self.config.convolution_bias: + return False + + # In transformers v5 (not v4), `convolution_bias=False` is + # propagated from parakeet config. If `False`, torch.conv1d will + # *skip registering the param*, thus it will be missing in the + # module's named params. *If* you happen to also have the bias + # tensors in the weights, it will cause a mismatch between the + # weights and the params. + # This allows us to have `convolution_bias=False` in the sound config, + # but still allow for the weights to exist. + + return target_name.endswith( + ( + ".conv.pointwise_conv1.bias", + ".conv.depthwise_conv.bias", + ".conv.pointwise_conv2.bias", + ) + ) + EPSILON = 1e-5 LOG_ZERO_GUARD_VALUE = 2**-24 From 79e799ebbd0a4472f8c6bd846ec676fded664138 Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Thu, 16 Apr 2026 19:26:55 -0400 Subject: [PATCH 072/150] [Bugfix] Temporarily disable B200 fp4 MoE layer tests (#40057) Signed-off-by: Bill Nell --- tests/kernels/moe/test_moe_layer.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index aa2948b8e98..01664d6a932 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -465,6 +465,14 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: if config.enable_eplb and config.ep_size == 1: return False, "EPLB only works with EP+DP" + # Disable fp4 tests until flashinfer is updated or the Dockerfile is + # modified to install cublasLt.h. See #39525. + if ( + config.quantization == "modelopt_fp4" + and current_platform.is_device_capability_family(100) + ): + return False, "Temporarily skip until #39525 is resolved" + return True, None From bf9a5ddb24af910f53e7d20305045010d7471072 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Thu, 16 Apr 2026 16:27:51 -0700 Subject: [PATCH 073/150] [MLA] Optimize mla indexer prepare uniform decode for MTP > 1 (#39458) Signed-off-by: Giancarlo Delfin --- vllm/v1/attention/backends/mla/indexer.py | 143 +++++++++++++++------- 1 file changed, 100 insertions(+), 43 deletions(-) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 402dfc0c74f..3b719d10ff8 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -8,6 +8,7 @@ import vllm.envs as envs from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton from vllm.utils.deep_gemm import ( get_paged_mqa_logits_metadata, has_deep_gemm, @@ -30,6 +31,40 @@ from vllm.v1.worker.cp_utils import get_total_cp_world_size logger = init_logger(__name__) +@triton.jit +def _prepare_uniform_decode_kernel( + seq_lens_ptr, + decode_seq_lens_ptr, + block_table_ptr, + block_table_stride, + expanded_block_table_ptr, + expanded_bt_stride, + decode_lens_ptr, + max_decode_len, + BLOCK_SIZE: tl.constexpr, +): + idx = tl.program_id(0) + req_id = idx // max_decode_len + local_idx = idx % max_decode_len + + # Compute number of KVs attended to by this token. + seq_len = tl.load(seq_lens_ptr + req_id) + per_token_seq_len = seq_len - max_decode_len + local_idx + 1 + tl.store(decode_seq_lens_ptr + idx, per_token_seq_len) + + # Copy block table row. + src = block_table_ptr + req_id * block_table_stride + dst = expanded_block_table_ptr + idx * expanded_bt_stride + for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE): + off = i + tl.arange(0, BLOCK_SIZE) + mask = off < expanded_bt_stride + src_block = tl.load(src + off, mask=mask) + tl.store(dst + off, src_block, mask=mask) + + # All reqs now have decode_len = 1. + tl.store(decode_lens_ptr + idx, 1) + + def split_indexer_prefill_chunks( seq_lens_cpu: torch.Tensor, query_lens_cpu: torch.Tensor, @@ -405,52 +440,75 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): Returns (seq_lens, block_table, decode_lens, batch_size, requires_padding). seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, next_n) for native MTP. """ + min_decode_len = int(decode_lens_cpu.min().item()) if not use_native and max_decode_len > 1: assert self.decode_seq_lens_buffer.dim() == 1 - # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is - # padding) and decode_lens [3, 1, 4, 0] in the below example comments. - # The context lengths are therefore - # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0]. - - # 3 + 1 + 4 + 0 = 8 - actual_expanded = int(decode_lens_cpu.sum().item()) - - # Fuse expanded_base and expanded_starts into a single repeat_interleave: - # seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1 - # where context_start[b] = seq_lens[b] - decode_lens[b]. - # Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8] - # expanded_offsets = [7, 7, 7, 3, 4, 4, 4, 4] - # result = [8, 9, 10, 7, 9, 10, 11, 12] - expanded_offsets = torch.repeat_interleave( - seq_lens - decode_lens - query_start_loc, - decode_lens, - output_size=actual_expanded, - ) - - # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space - self.decode_seq_lens_buffer[:actual_expanded] = ( - expanded_offsets + self.arange_buffer[:actual_expanded] + 1 - ) - self.decode_seq_lens_buffer[actual_expanded:] = 0 - seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens] - - # Give each of the flattened entries the same block table row as the - # original request. - self.expanded_block_table_buffer[:actual_expanded] = ( - torch.repeat_interleave( - block_table, decode_lens, dim=0, output_size=actual_expanded + if min_decode_len == max_decode_len: + # Uniform decode lengths. + num_decode_tokens = num_decodes * max_decode_len + _prepare_uniform_decode_kernel[(num_decode_tokens,)]( + seq_lens, + self.decode_seq_lens_buffer, + block_table, + block_table.stride(0), + self.expanded_block_table_buffer, + self.expanded_block_table_buffer.stride(0), + self.decode_lens_buffer, + max_decode_len, + BLOCK_SIZE=1024, ) - ) - if actual_expanded < num_decode_tokens: - self.expanded_block_table_buffer[ - actual_expanded:num_decode_tokens, 0 - ] = 0 - block_table = self.expanded_block_table_buffer[:num_decode_tokens] + self.decode_seq_lens_buffer[num_decode_tokens:] = 0 + seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens] + block_table = self.expanded_block_table_buffer[:num_decode_tokens] + decode_lens = self.decode_lens_buffer[:num_decode_tokens] + return seq_lens, block_table, decode_lens, num_decode_tokens, False + else: + # Variable decode lengths. + # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is + # padding) and decode_lens [3, 1, 4, 0] in the below example comments. + # The context lengths are therefore + # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0]. - # All reqs now have decode_len=1 - self.decode_lens_buffer[:num_decode_tokens] = 1 - decode_lens = self.decode_lens_buffer[:num_decode_tokens] - return seq_lens, block_table, decode_lens, num_decode_tokens, False + # 3 + 1 + 4 + 0 = 8 + actual_expanded = int(decode_lens_cpu.sum().item()) + + # Fuse expanded_base and expanded_starts into a single + # repeat_interleave: + # seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1 + # where context_start[b] = seq_lens[b] - decode_lens[b]. + # Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8] + # expanded_offsets = [7, 7, 7, 3, 4, 4, 4, 4] + # result = [8, 9, 10, 7, 9, 10, 11, 12] + expanded_offsets = torch.repeat_interleave( + seq_lens - decode_lens - query_start_loc, + decode_lens, + output_size=actual_expanded, + ) + + # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space + self.decode_seq_lens_buffer[:actual_expanded] = ( + expanded_offsets + self.arange_buffer[:actual_expanded] + 1 + ) + self.decode_seq_lens_buffer[actual_expanded:] = 0 + seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens] + + # Give each of the flattened entries the same block table row as the + # original request. + self.expanded_block_table_buffer[:actual_expanded] = ( + torch.repeat_interleave( + block_table, decode_lens, dim=0, output_size=actual_expanded + ) + ) + if actual_expanded < num_decode_tokens: + self.expanded_block_table_buffer[ + actual_expanded:num_decode_tokens, 0 + ] = 0 + block_table = self.expanded_block_table_buffer[:num_decode_tokens] + + # All reqs now have decode_len=1 + self.decode_lens_buffer[:num_decode_tokens] = 1 + decode_lens = self.decode_lens_buffer[:num_decode_tokens] + return seq_lens, block_table, decode_lens, num_decode_tokens, False else: # Native path: plain decode (next_n==1) or spec decode # with 2D per-token context lengths (next_n > 1). @@ -459,7 +517,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): # decode_len < next_n due to padding or short prefills), the simple # reshape in sparse_attn_indexer won't work. Use pack_seq_triton # (requires_padding) instead. - min_decode_len = int(decode_lens_cpu.min().item()) requires_padding = min_decode_len != max_decode_len if use_native and next_n > 1: assert self.decode_seq_lens_buffer.dim() == 2 From 4c47710bf707e93bdf83391948f844e461b7ba4f Mon Sep 17 00:00:00 2001 From: Shinichi Hemmi <50256998+Alnusjaponica@users.noreply.github.com> Date: Fri, 17 Apr 2026 09:54:32 +0900 Subject: [PATCH 074/150] [CI/Build] Apply ruff formatter to pass pre-commit (#40078) Signed-off-by: Hemmi Shinichi --- .../passes/test_vllm_fusion_pattern_matcher_pass.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py b/tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py index ee2eef009fc..381c215a9ae 100644 --- a/tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py +++ b/tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py @@ -6,14 +6,13 @@ import torch import vllm.config from tests.compile.backend import TestBackend -from vllm.platforms import current_platform from vllm.compilation.passes.vllm_inductor_pass import ( VllmFusionPatternMatcherPass, VllmPatternMatcherPass, VllmPatternReplacement, ) from vllm.config import CompilationConfig, CompilationMode, VllmConfig - +from vllm.platforms import current_platform class ReluToAbsPattern(VllmPatternReplacement): @@ -58,7 +57,6 @@ class ExpToSqrtPattern(VllmPatternReplacement): return [self.empty_fp32(4)] - class ReluFusionPass(VllmFusionPatternMatcherPass): def __init__(self, config: VllmConfig) -> None: super().__init__(config, "test_relu_fusion") @@ -72,13 +70,13 @@ class TwoPatternFusionPass(VllmFusionPatternMatcherPass): self.register(ExpToSqrtPattern()) - @pytest.fixture def vllm_config(): return VllmConfig( compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE), ) + @pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA") def test_register_tracks_patterns(vllm_config): """register() appends each VllmPatternReplacement to _pattern_replacements.""" @@ -96,7 +94,7 @@ def test_uuid_stable(vllm_config): with vllm.config.set_current_vllm_config(vllm_config): p1 = ReluFusionPass(vllm_config) p2 = ReluFusionPass(vllm_config) - p3= TwoPatternFusionPass(vllm_config) + p3 = TwoPatternFusionPass(vllm_config) assert p1.uuid() == p2.uuid() assert p1.uuid() != p3.uuid() From 1948d0c467da0b273b287665028ae82b9441c7c7 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Thu, 16 Apr 2026 22:48:37 -0400 Subject: [PATCH 075/150] [UX] Defer some imports on CLI paths to save ~2s (#40056) Signed-off-by: mgoin --- vllm/benchmarks/sweep/plot.py | 30 +++++++++-------- vllm/benchmarks/sweep/plot_pareto.py | 30 +++++++++-------- vllm/entrypoints/cli/__init__.py | 17 ---------- vllm/entrypoints/cli/benchmark/main.py | 46 +++++++++++++++++++------- vllm/env_override.py | 5 ++- 5 files changed, 68 insertions(+), 60 deletions(-) diff --git a/vllm/benchmarks/sweep/plot.py b/vllm/benchmarks/sweep/plot.py index 156e18f697f..2d369280444 100644 --- a/vllm/benchmarks/sweep/plot.py +++ b/vllm/benchmarks/sweep/plot.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from functools import partial from pathlib import Path from types import TracebackType -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar from typing_extensions import Self, override @@ -17,20 +17,8 @@ from vllm.utils.import_utils import PlaceholderModule from .utils import sanitize_filename -try: - import matplotlib.pyplot as plt -except ImportError: - plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot") - -try: +if TYPE_CHECKING: import pandas as pd -except ImportError: - pd = PlaceholderModule("pandas") - -try: - import seaborn as sns -except ImportError: - seaborn = PlaceholderModule("seaborn") @dataclass @@ -265,6 +253,20 @@ def _plot_fig( fig_height: float, fig_dpi: int, ): + # Lazy-import matplotlib/pandas/seaborn + try: + import matplotlib.pyplot as plt + except ImportError: + plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot") + try: + import pandas as pd + except ImportError: + pd = PlaceholderModule("pandas") + try: + import seaborn as sns + except ImportError: + sns = PlaceholderModule("seaborn") + fig_group, fig_data = fig_group_data row_groups = full_groupby( diff --git a/vllm/benchmarks/sweep/plot_pareto.py b/vllm/benchmarks/sweep/plot_pareto.py index 365e87f757d..8ec309a7a10 100644 --- a/vllm/benchmarks/sweep/plot_pareto.py +++ b/vllm/benchmarks/sweep/plot_pareto.py @@ -6,7 +6,7 @@ from concurrent.futures import ProcessPoolExecutor from dataclasses import dataclass from functools import partial from pathlib import Path -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar from vllm.utils.collection_utils import full_groupby from vllm.utils.import_utils import PlaceholderModule @@ -14,20 +14,8 @@ from vllm.utils.import_utils import PlaceholderModule from .plot import DummyExecutor, _json_load_bytes from .utils import sanitize_filename -try: - import matplotlib.pyplot as plt -except ImportError: - plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot") - -try: +if TYPE_CHECKING: import pandas as pd -except ImportError: - pd = PlaceholderModule("pandas") - -try: - import seaborn as sns -except ImportError: - seaborn = PlaceholderModule("seaborn") def _first_present(run_data: dict[str, object], keys: list[str]): @@ -195,6 +183,20 @@ def _plot_fig( print("[END FIGURE]") return + # Lazy-import matplotlib/pandas/seaborn + try: + import matplotlib.pyplot as plt + except ImportError: + plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot") + try: + import pandas as pd + except ImportError: + pd = PlaceholderModule("pandas") + try: + import seaborn as sns + except ImportError: + sns = PlaceholderModule("seaborn") + df = pd.DataFrame.from_records(fig_data) df = df.dropna(subset=["tokens_per_user", "tokens_per_gpu"]) diff --git a/vllm/entrypoints/cli/__init__.py b/vllm/entrypoints/cli/__init__.py index 704d94d36f7..208f01a7cb5 100644 --- a/vllm/entrypoints/cli/__init__.py +++ b/vllm/entrypoints/cli/__init__.py @@ -1,19 +1,2 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.entrypoints.cli.benchmark.latency import BenchmarkLatencySubcommand -from vllm.entrypoints.cli.benchmark.mm_processor import ( - BenchmarkMMProcessorSubcommand, -) -from vllm.entrypoints.cli.benchmark.serve import BenchmarkServingSubcommand -from vllm.entrypoints.cli.benchmark.startup import BenchmarkStartupSubcommand -from vllm.entrypoints.cli.benchmark.sweep import BenchmarkSweepSubcommand -from vllm.entrypoints.cli.benchmark.throughput import BenchmarkThroughputSubcommand - -__all__: list[str] = [ - "BenchmarkLatencySubcommand", - "BenchmarkMMProcessorSubcommand", - "BenchmarkServingSubcommand", - "BenchmarkStartupSubcommand", - "BenchmarkSweepSubcommand", - "BenchmarkThroughputSubcommand", -] diff --git a/vllm/entrypoints/cli/benchmark/main.py b/vllm/entrypoints/cli/benchmark/main.py index 48f34fce1d4..f64de4cf673 100644 --- a/vllm/entrypoints/cli/benchmark/main.py +++ b/vllm/entrypoints/cli/benchmark/main.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import sys import typing from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase @@ -14,6 +15,17 @@ else: FlexibleArgumentParser = argparse.ArgumentParser +def _import_bench_subcommand_modules() -> None: + # Imported lazily so `BenchmarkSubcommandBase` subclasses register only + # when `vllm bench` is actually invoked. + import vllm.entrypoints.cli.benchmark.latency # noqa: F401 + import vllm.entrypoints.cli.benchmark.mm_processor # noqa: F401 + import vllm.entrypoints.cli.benchmark.serve # noqa: F401 + import vllm.entrypoints.cli.benchmark.startup # noqa: F401 + import vllm.entrypoints.cli.benchmark.sweep # noqa: F401 + import vllm.entrypoints.cli.benchmark.throughput # noqa: F401 + + class BenchmarkSubcommand(CLISubcommand): """The `bench` subcommand for the vLLM CLI.""" @@ -38,18 +50,28 @@ class BenchmarkSubcommand(CLISubcommand): ) bench_subparsers = bench_parser.add_subparsers(required=True, dest="bench_type") - for cmd_cls in BenchmarkSubcommandBase.__subclasses__(): - cmd_subparser = bench_subparsers.add_parser( - cmd_cls.name, - help=cmd_cls.help, - description=cmd_cls.help, - usage=f"vllm {self.name} {cmd_cls.name} [options]", - ) - cmd_subparser.set_defaults(dispatch_function=cmd_cls.cmd) - cmd_cls.add_cli_args(cmd_subparser) - cmd_subparser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format( - subcmd=f"{self.name} {cmd_cls.name}" - ) + # Only build the nested bench subparsers when the user is actually + # invoking `bench`; otherwise we'd drag in imports + # unnecessarily on every `vllm --help` and `vllm serve`. + # Scan for the first positional arg so global flags (e.g. `-v`) + # before the subcommand don't break detection. + first_positional = next( + (arg for arg in sys.argv[1:] if not arg.startswith("-")), None + ) + if first_positional == self.name: + _import_bench_subcommand_modules() + for cmd_cls in BenchmarkSubcommandBase.__subclasses__(): + cmd_subparser = bench_subparsers.add_parser( + cmd_cls.name, + help=cmd_cls.help, + description=cmd_cls.help, + usage=f"vllm {self.name} {cmd_cls.name} [options]", + ) + cmd_subparser.set_defaults(dispatch_function=cmd_cls.cmd) + cmd_cls.add_cli_args(cmd_subparser) + cmd_subparser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format( + subcmd=f"{self.name} {cmd_cls.name}" + ) return bench_parser diff --git a/vllm/env_override.py b/vllm/env_override.py index 2e3f02866ab..a19b6e25541 100644 --- a/vllm/env_override.py +++ b/vllm/env_override.py @@ -100,10 +100,9 @@ logger = init_logger(__name__) # it avoids unintentional cuda initialization from torch.cuda.is_available() os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1" -# see https://github.com/vllm-project/vllm/issues/10480 +# see https://github.com/vllm-project/vllm/issues/10480 and +# https://github.com/vllm-project/vllm/issues/10619. os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1" -# see https://github.com/vllm-project/vllm/issues/10619 -torch._inductor.config.compile_threads = 1 # Enable Triton autotuning result caching to disk by default. # Without this, Triton re-runs autotuning on every process restart, From 978a4462bbc529ff204647543526e4caa08ed974 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Fri, 17 Apr 2026 12:17:39 +0800 Subject: [PATCH 076/150] [CI Failure] Fix Plugin Tests (2 GPUs) Failure (#40083) Signed-off-by: wang.yuqi Co-authored-by: Michele Gazzetti --- vllm/entrypoints/pooling/pooling/protocol.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/entrypoints/pooling/pooling/protocol.py b/vllm/entrypoints/pooling/pooling/protocol.py index b75c34e9790..ab9b132fee3 100644 --- a/vllm/entrypoints/pooling/pooling/protocol.py +++ b/vllm/entrypoints/pooling/pooling/protocol.py @@ -97,6 +97,11 @@ class IOProcessorRequest(PoolingBasicRequestMixin, EncodingRequestMixin, Generic max_total_tokens_param="max_model_len", ) + def to_pooling_params(self): + return PoolingParams( + task=self.task, + ) + class IOProcessorResponse(OpenAIBaseModel, Generic[T]): request_id: str | None = None From bf45e6d0a558da2b8d7b60efb07b4aa394f3b60b Mon Sep 17 00:00:00 2001 From: z1ying <55220715+z1ying@users.noreply.github.com> Date: Thu, 16 Apr 2026 22:42:52 -0700 Subject: [PATCH 077/150] [Doc] Add Gemma 4 to supported models list (#39607) Signed-off-by: z1ying Signed-off-by: Ziying Tao --- docs/models/supported_models.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index ef1f5901ed0..f0eb3781fba 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -400,6 +400,7 @@ th { | `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ | | `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ | | `Gemma3nForCausalLM` | Gemma 3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | +| `Gemma4ForCausalLM` | Gemma 4 | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ | | `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ | | `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ | | `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ | @@ -554,6 +555,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ | | `Gemma3ForConditionalGeneration` | Gemma 3 | T + IE+ | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ | | `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | +| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I+ + V + A* | `google/gemma-4-E2B-it`, etc. | | ✅︎ | | `GLM4VForCausalLM`^ | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ | | `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + IE+ + VE+ | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ | | `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + IE+ + VE+ | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ | @@ -633,6 +635,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor ^ You need to set the architecture name via `--hf-overrides` to match the one in vLLM.
E Pre-computed embeddings can be inputted for this modality.
+ Multiple items can be inputted per text prompt for this modality. +* Only specific variants of the model support this modality (see notes below).
!!! note `Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its @@ -643,6 +646,11 @@ Some models are supported only via the [Transformers modeling backend](#transfor - Both audio and vision MM encoders use `transformers.AutoModel` implementation. - There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups. +!!! note + For `Gemma4ForConditionalGeneration`: + - audio input is only supported by the `gemma-4-E2B` and `gemma-4-E4B` variants. + - The model does not ingest videos directly. However, vLLM’s Gemma 4 implementation supports video inputs by handling video processing internally. Users can send videos directly in the message structure to vLLM, where they are converted into text and image frames before being passed to the model. + !!! note For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently. From 4f436782afd0b21d6754ea6bc4b80639f737bbc1 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Fri, 17 Apr 2026 16:56:22 +0800 Subject: [PATCH 078/150] [Misc] Improve new PR bot trigger condition (#40114) Signed-off-by: DarkLight1337 --- .github/workflows/new_pr_bot.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/new_pr_bot.yml b/.github/workflows/new_pr_bot.yml index ef5e30952c6..27100f9f4da 100644 --- a/.github/workflows/new_pr_bot.yml +++ b/.github/workflows/new_pr_bot.yml @@ -62,14 +62,14 @@ jobs: const prAuthor = context.payload.pull_request.user.login; const { data: searchResults } = await github.rest.search.issuesAndPullRequests({ - q: `repo:${owner}/${repo} type:pr author:${prAuthor}`, + q: `repo:${owner}/${repo} type:pr is:merged author:${prAuthor}`, per_page: 1, }); - const authorPRCount = searchResults.total_count; - console.log(`Found ${authorPRCount} PRs by ${prAuthor}`); + const mergedPRCount = searchResults.total_count; + console.log(`Found ${mergedPRCount} merged PRs by ${prAuthor}`); - if (authorPRCount === 1) { + if (mergedPRCount === 0) { console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`); await github.rest.issues.createComment({ owner, @@ -98,5 +98,5 @@ jobs: ].join('\n'), }); } else { - console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`); + console.log(`Skipping comment for ${prAuthor} - not a first-time contributor (${mergedPRCount} merged PRs)`); } From 8d2cff8140eb2c5f6a7b28948c0944d11898d9d2 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Fri, 17 Apr 2026 18:13:31 +0800 Subject: [PATCH 079/150] [Examples] Resettle Observability examples. (#40123) Signed-off-by: wang.yuqi --- docs/design/metrics.md | 4 ++-- .../{online_serving => observability}/dashboards/README.md | 4 ++-- .../dashboards/grafana/README.md | 0 .../dashboards/grafana/performance_statistics.json | 0 .../dashboards/grafana/query_statistics.json | 0 .../dashboards/perses/README.md | 0 .../dashboards/perses/performance_statistics.yaml | 0 .../dashboards/perses/query_statistics.yaml | 0 .../metrics.py => observability/metrics/offline.py} | 0 .../{online_serving => observability}/opentelemetry/README.md | 0 .../opentelemetry/dummy_client.py | 0 .../prometheus_grafana/README.md | 0 .../prometheus_grafana/docker-compose.yaml | 0 .../prometheus_grafana/grafana.json | 0 .../prometheus_grafana/prometheus.yaml | 0 15 files changed, 4 insertions(+), 4 deletions(-) rename examples/{online_serving => observability}/dashboards/README.md (92%) rename examples/{online_serving => observability}/dashboards/grafana/README.md (100%) rename examples/{online_serving => observability}/dashboards/grafana/performance_statistics.json (100%) rename examples/{online_serving => observability}/dashboards/grafana/query_statistics.json (100%) rename examples/{online_serving => observability}/dashboards/perses/README.md (100%) rename examples/{online_serving => observability}/dashboards/perses/performance_statistics.yaml (100%) rename examples/{online_serving => observability}/dashboards/perses/query_statistics.yaml (100%) rename examples/{offline_inference/metrics.py => observability/metrics/offline.py} (100%) rename examples/{online_serving => observability}/opentelemetry/README.md (100%) rename examples/{online_serving => observability}/opentelemetry/dummy_client.py (100%) rename examples/{online_serving => observability}/prometheus_grafana/README.md (100%) rename examples/{online_serving => observability}/prometheus_grafana/docker-compose.yaml (100%) rename examples/{online_serving => observability}/prometheus_grafana/grafana.json (100%) rename examples/{online_serving => observability}/prometheus_grafana/prometheus.yaml (100%) diff --git a/docs/design/metrics.md b/docs/design/metrics.md index be917c0dc61..0ae42039976 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -42,7 +42,7 @@ These are documented under [Inferencing and Serving -> Production Metrics](../us ### Grafana Dashboard -vLLM also provides [a reference example](../../examples/online_serving/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard. +vLLM also provides [a reference example](../../examples/observability/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard. The subset of metrics exposed in the Grafana dashboard gives us an indication of which metrics are especially important: @@ -657,7 +657,7 @@ vLLM has support for OpenTelemetry tracing: - Added by and reinstated by - Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces` - [OpenTelemetry blog post](https://opentelemetry.io/blog/2024/llm-observability/) -- [User-facing docs](../../examples/online_serving/opentelemetry/README.md) +- [User-facing docs](../../examples/observability/opentelemetry/README.md) - [Blog post](https://medium.com/@ronen.schaffer/follow-the-trail-supercharging-vllm-with-opentelemetry-distributed-tracing-aa655229b46f) - [IBM product docs](https://www.ibm.com/docs/en/instana-observability/current?topic=mgaa-monitoring-large-language-models-llms-vllm-public-preview) diff --git a/examples/online_serving/dashboards/README.md b/examples/observability/dashboards/README.md similarity index 92% rename from examples/online_serving/dashboards/README.md rename to examples/observability/dashboards/README.md index 10b9a864f57..e5f5010a42a 100644 --- a/examples/online_serving/dashboards/README.md +++ b/examples/observability/dashboards/README.md @@ -74,8 +74,8 @@ percli apply -f perses/performance_statistics.yaml For detailed deployment instructions and platform-specific options, see: -- **[Grafana Documentation](./grafana)** - JSON dashboards, operator usage, manual import -- **[Perses Documentation](./perses)** - YAML specs, CLI usage, operator wrapping +- **[Grafana Documentation](grafana)** - JSON dashboards, operator usage, manual import +- **[Perses Documentation](perses)** - YAML specs, CLI usage, operator wrapping ## Contributing diff --git a/examples/online_serving/dashboards/grafana/README.md b/examples/observability/dashboards/grafana/README.md similarity index 100% rename from examples/online_serving/dashboards/grafana/README.md rename to examples/observability/dashboards/grafana/README.md diff --git a/examples/online_serving/dashboards/grafana/performance_statistics.json b/examples/observability/dashboards/grafana/performance_statistics.json similarity index 100% rename from examples/online_serving/dashboards/grafana/performance_statistics.json rename to examples/observability/dashboards/grafana/performance_statistics.json diff --git a/examples/online_serving/dashboards/grafana/query_statistics.json b/examples/observability/dashboards/grafana/query_statistics.json similarity index 100% rename from examples/online_serving/dashboards/grafana/query_statistics.json rename to examples/observability/dashboards/grafana/query_statistics.json diff --git a/examples/online_serving/dashboards/perses/README.md b/examples/observability/dashboards/perses/README.md similarity index 100% rename from examples/online_serving/dashboards/perses/README.md rename to examples/observability/dashboards/perses/README.md diff --git a/examples/online_serving/dashboards/perses/performance_statistics.yaml b/examples/observability/dashboards/perses/performance_statistics.yaml similarity index 100% rename from examples/online_serving/dashboards/perses/performance_statistics.yaml rename to examples/observability/dashboards/perses/performance_statistics.yaml diff --git a/examples/online_serving/dashboards/perses/query_statistics.yaml b/examples/observability/dashboards/perses/query_statistics.yaml similarity index 100% rename from examples/online_serving/dashboards/perses/query_statistics.yaml rename to examples/observability/dashboards/perses/query_statistics.yaml diff --git a/examples/offline_inference/metrics.py b/examples/observability/metrics/offline.py similarity index 100% rename from examples/offline_inference/metrics.py rename to examples/observability/metrics/offline.py diff --git a/examples/online_serving/opentelemetry/README.md b/examples/observability/opentelemetry/README.md similarity index 100% rename from examples/online_serving/opentelemetry/README.md rename to examples/observability/opentelemetry/README.md diff --git a/examples/online_serving/opentelemetry/dummy_client.py b/examples/observability/opentelemetry/dummy_client.py similarity index 100% rename from examples/online_serving/opentelemetry/dummy_client.py rename to examples/observability/opentelemetry/dummy_client.py diff --git a/examples/online_serving/prometheus_grafana/README.md b/examples/observability/prometheus_grafana/README.md similarity index 100% rename from examples/online_serving/prometheus_grafana/README.md rename to examples/observability/prometheus_grafana/README.md diff --git a/examples/online_serving/prometheus_grafana/docker-compose.yaml b/examples/observability/prometheus_grafana/docker-compose.yaml similarity index 100% rename from examples/online_serving/prometheus_grafana/docker-compose.yaml rename to examples/observability/prometheus_grafana/docker-compose.yaml diff --git a/examples/online_serving/prometheus_grafana/grafana.json b/examples/observability/prometheus_grafana/grafana.json similarity index 100% rename from examples/online_serving/prometheus_grafana/grafana.json rename to examples/observability/prometheus_grafana/grafana.json diff --git a/examples/online_serving/prometheus_grafana/prometheus.yaml b/examples/observability/prometheus_grafana/prometheus.yaml similarity index 100% rename from examples/online_serving/prometheus_grafana/prometheus.yaml rename to examples/observability/prometheus_grafana/prometheus.yaml From c0c98b8b9a392c7e8b36b68cf477e245dda48d80 Mon Sep 17 00:00:00 2001 From: Maral Date: Fri, 17 Apr 2026 18:20:32 +0800 Subject: [PATCH 080/150] [Bugfix] Add Marlin kernel in block scaled mm kernel selection. (#40105) Signed-off-by: maral --- vllm/model_executor/kernels/linear/__init__.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 6cbb65e26d6..5d513f767f0 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -186,12 +186,13 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = # in priority/performance order (when available) _POSSIBLE_FP8_BLOCK_KERNELS: dict[ - PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel]] + PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]] ] = { PlatformEnum.CUDA: [ FlashInferFp8DeepGEMMDynamicBlockScaledKernel, DeepGemmFp8BlockScaledMMKernel, CutlassFp8BlockScaledMMKernel, + MarlinFP8ScaledMMLinearKernel, TritonFp8BlockScaledMMKernel, ], PlatformEnum.ROCM: [ @@ -392,6 +393,19 @@ def init_fp8_linear_kernel( scope="global", ) + # TODO make scaled_mm kernels inherit from MMLinearKernel + # only MarlinFP8ScaledMMLinearKernel is a type of FP8ScaledMMLinearKernel. + if issubclass(kernel_type, FP8ScaledMMLinearKernel): + return kernel_type( + scaled_mm_linear_kernel_config, + layer_param_names=[ + "weight", + "weight_scale", + "input_scale", + "input_scale_ub", + ], + ) + return kernel_type( scaled_mm_linear_kernel_config, ) @@ -399,7 +413,7 @@ def init_fp8_linear_kernel( else: kernel_type = choose_scaled_mm_linear_kernel( config=scaled_mm_linear_kernel_config, - possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[misc] + possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type] force_kernel=force_kernel, ) if module_name: From 79a5b6325396683d3063c41fd1a140c56c89e982 Mon Sep 17 00:00:00 2001 From: Or Ozeri Date: Fri, 17 Apr 2026 15:13:55 +0300 Subject: [PATCH 081/150] [kv_offload]: Fix num CPU blocks for UniformTypeKVCacheSpecs (#39617) Signed-off-by: Or Ozeri --- vllm/v1/kv_offload/cpu/spec.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index a0fafdeb118..8d4b744a0b4 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -26,17 +26,13 @@ class CPUOffloadingSpec(OffloadingSpec): # calculate kv_bytes_per_offloaded_block assert kv_cache_config is not None - page_sizes = { - kv_cache_group.kv_cache_spec.page_size_bytes - for kv_cache_group in kv_cache_config.kv_cache_groups - } - assert len(page_sizes) == 1 - page_size_bytes = page_sizes.pop() - kv_bytes_per_block = ( - page_size_bytes - * len(kv_cache_config.kv_cache_tensors) - * vllm_config.parallel_config.world_size - ) + if kv_cache_config.num_blocks > 0: + total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors) + kv_bytes_per_block = ( + total_gpu_kv_bytes // kv_cache_config.num_blocks + ) * vllm_config.parallel_config.world_size + else: + kv_bytes_per_block = 0 kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor self.num_blocks = ( From b1dc87a0989fd98a614e4e7e6b318a19e8c5c6f9 Mon Sep 17 00:00:00 2001 From: Lukas Geiger Date: Fri, 17 Apr 2026 13:37:21 +0100 Subject: [PATCH 082/150] [Models][Gemma4] Prevent GPU/CPU sync in `embed_input_ids` (#39234) Signed-off-by: Lukas Geiger --- vllm/model_executor/models/gemma4_mm.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index cd0ca37616b..08978ec3f44 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -1254,9 +1254,10 @@ class Gemma4ForConditionalGeneration( # computation (using token_type_ids == 0 as text_mask). # Replicate this: map image token positions to token 0. if is_multimodal is not None: - is_multimodal = is_multimodal.to(input_ids.device) ple_input_ids = torch.where( - is_multimodal, torch.zeros_like(input_ids), input_ids + is_multimodal.to(input_ids.device, non_blocking=True), + torch.zeros_like(input_ids), + input_ids, ) else: ple_input_ids = input_ids From d02421a7dbd85eb173cb2620da3dbc16d81135f4 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Fri, 17 Apr 2026 21:01:08 +0800 Subject: [PATCH 083/150] [CPU] Refactor CPU affinity and memory management (#39781) Signed-off-by: jiang1.li --- .../run-cpu-distributed-smoke-test.sh | 38 +- .github/workflows/macos-smoke-test.yml | 1 + cmake/cpu_extension.cmake | 29 +- csrc/cpu/torch_bindings.cpp | 4 + csrc/cpu/utils.cpp | 79 ++- docker/Dockerfile.cpu | 3 +- .../models/language/generation/test_common.py | 7 +- vllm/config/cache.py | 3 - vllm/platforms/cpu.py | 208 ++------ vllm/utils/cpu_resource_utils.py | 173 +++++++ vllm/utils/ompmultiprocessing.py | 454 +++++++++++------- vllm/v1/executor/multiproc_executor.py | 21 +- vllm/v1/worker/cpu_model_runner.py | 16 +- vllm/v1/worker/cpu_worker.py | 108 ++++- 14 files changed, 721 insertions(+), 423 deletions(-) create mode 100644 vllm/utils/cpu_resource_utils.py diff --git a/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh index 4ff15067aa0..f12bb524d4c 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh @@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then exit 1 fi -#echo "--- DP+TP" -#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 & -#server_pid=$! -#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1 -#vllm bench serve \ -# --backend vllm \ -# --dataset-name random \ -# --model meta-llama/Llama-3.2-3B-Instruct \ -# --num-prompts 20 \ -# --result-dir ./test_results \ -# --result-filename dp_pp.json \ -# --save-result \ -# --endpoint /v1/completions -#kill -s SIGTERM $server_pid; wait $server_pid || true -#failed_req=$(jq '.failed' ./test_results/dp_pp.json) -#if [ "$failed_req" -ne 0 ]; then -# echo "Some requests were failed!" -# exit 1 -#fi +echo "--- DP+TP" +vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 & +server_pid=$! +timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1 +vllm bench serve \ + --backend vllm \ + --dataset-name random \ + --model meta-llama/Llama-3.2-3B-Instruct \ + --num-prompts 20 \ + --result-dir ./test_results \ + --result-filename dp_pp.json \ + --save-result \ + --endpoint /v1/completions +kill -s SIGTERM $server_pid; wait $server_pid || true +failed_req=$(jq '.failed' ./test_results/dp_pp.json) +if [ "$failed_req" -ne 0 ]; then + echo "Some requests were failed!" + exit 1 +fi diff --git a/.github/workflows/macos-smoke-test.yml b/.github/workflows/macos-smoke-test.yml index b8e7cebf2dc..ea1c8b0feac 100644 --- a/.github/workflows/macos-smoke-test.yml +++ b/.github/workflows/macos-smoke-test.yml @@ -45,6 +45,7 @@ jobs: - name: Smoke test vllm serve run: | # Start server in background + VLLM_CPU_KVCACHE_SPACE=1 \ vllm serve Qwen/Qwen3-0.6B \ --max-model-len=2K \ --load-format=dummy \ diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index f45280ee488..698177d0756 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -30,6 +30,21 @@ else() list(APPEND CXX_COMPILE_FLAGS "-fopenmp" "-DVLLM_CPU_EXTENSION") + + # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0) + # and create a local shim dir with it + vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR) + + find_library(OPEN_MP + NAMES gomp + PATHS ${VLLM_TORCH_GOMP_SHIM_DIR} + NO_DEFAULT_PATH + REQUIRED + ) + # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch + if (OPEN_MP) + set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}") + endif() endif() if (NOT MACOSX_FOUND) @@ -175,20 +190,6 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND if(NOT NPROC) set(NPROC 4) endif() - # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0) - # and create a local shim dir with it - vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR) - - find_library(OPEN_MP - NAMES gomp - PATHS ${VLLM_TORCH_GOMP_SHIM_DIR} - NO_DEFAULT_PATH - REQUIRED - ) - # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch - if (OPEN_MP) - set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}") - endif() # Fetch and populate ACL if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}") diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 83bb57c9c09..b2b85190eec 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -141,6 +141,8 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc, torch::Tensor slot_mapping, const int64_t block_size); +void init_cpu_memory_env(std::vector node_ids); + namespace cpu_utils { void eagle_prepare_inputs_padded_kernel_impl( const torch::Tensor& cu_num_draft_tokens, @@ -431,6 +433,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "block_size) -> ()", &compute_slot_mapping_kernel_impl); + ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env); + // Speculative decoding kernels ops.def( "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, " diff --git a/csrc/cpu/utils.cpp b/csrc/cpu/utils.cpp index 55d0a3d1757..fd06e01e4d6 100644 --- a/csrc/cpu/utils.cpp +++ b/csrc/cpu/utils.cpp @@ -13,13 +13,80 @@ #include "cpu/utils.hpp" #ifdef VLLM_NUMA_DISABLED -std::string init_cpu_threads_env(const std::string& cpu_ids) { - return std::string( - "Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has " - "no effect to setup thread affinity."); -} +void init_cpu_memory_env(std::vector node_ids) {} +#else +void init_cpu_memory_env(std::vector node_ids) { + // Memory node binding + if (numa_available() != -1) { + // Concatenate all node_ids into a single comma-separated string + if (!node_ids.empty()) { + std::string node_ids_str; + for (const int node_id : node_ids) { + if (!node_ids_str.empty()) { + node_ids_str += ","; + } + node_ids_str += std::to_string(node_id); + } -#endif + bitmask* mask = numa_parse_nodestring(node_ids_str.c_str()); + bitmask* src_mask = numa_get_mems_allowed(); + + int pid = getpid(); + + if (mask && src_mask) { + // move all existing pages to the specified numa node. + *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp); + int page_num = numa_migrate_pages(pid, src_mask, mask); + if (page_num == -1) { + TORCH_WARN("numa_migrate_pages failed. errno: " + + std::to_string(errno)); + } + + // Restrict memory allocation to the selected NUMA node(s). + // Enhances memory locality for the threads bound to those NUMA CPUs. + if (node_ids.size() > 1) { + errno = 0; + numa_set_interleave_mask(mask); + if (errno != 0) { + TORCH_WARN("numa_set_interleave_mask failed. errno: " + + std::to_string(errno)); + } else { + TORCH_WARN( + "NUMA binding: Using INTERLEAVE policy for memory " + "allocation across multiple NUMA nodes (nodes: " + + node_ids_str + + "). Memory allocations will be " + "interleaved across the specified NUMA nodes."); + } + } else { + errno = 0; + numa_set_membind(mask); + if (errno != 0) { + TORCH_WARN("numa_set_membind failed. errno: " + + std::to_string(errno)); + } else { + TORCH_WARN( + "NUMA binding: Using MEMBIND policy for memory " + "allocation on the NUMA nodes (" + + node_ids_str + + "). Memory allocations will be " + "strictly bound to these NUMA nodes."); + } + } + + numa_set_strict(1); + + numa_free_nodemask(mask); + numa_free_nodemask(src_mask); + } else { + TORCH_WARN( + "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " + + std::to_string(errno)); + } + } + } +} +#endif // VLLM_NUMA_DISABLED namespace cpu_utils { ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) { diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 77b449625dd..72ee002d90a 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -173,7 +173,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \ COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/dev.txt && \ + uv pip install -r requirements/lint.txt && \ + uv pip install -r requirements/test/cpu.txt && \ pre-commit install --hook-type pre-commit --hook-type commit-msg ENTRYPOINT ["bash"] diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py index ebf847b706c..6e5f1e32843 100644 --- a/tests/models/language/generation/test_common.py +++ b/tests/models/language/generation/test_common.py @@ -46,7 +46,7 @@ AITER_MODEL_LIST = [ ), pytest.param( "openai-community/gpt2", # gpt2 - marks=[pytest.mark.core_model, pytest.mark.cpu_model], + marks=[pytest.mark.core_model], ), pytest.param("Milos/slovak-gpt-j-405M"), # gptj pytest.param("bigcode/tiny_starcoder_py"), # gpt_bigcode @@ -143,11 +143,6 @@ def test_models( # in parts of the operators pytest.skip(f"Skipping '{model}' model test with AITER kernel.") - if current_platform.is_cpu() and model in ("openai-community/gpt2",): - # These models are sensitive to the rounding error - # Fuse ops to reduce rounding - monkeypatch.setenv("VLLM_CPU_CI_ENV", "0") - with hf_runner(model) as hf_model: hf_outputs = hf_model.generate_greedy_logprobs_limit( example_prompts, max_tokens, num_logprobs diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 47a655f22d5..b84df7773c6 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -101,8 +101,6 @@ class CacheConfig: kv_cache_dtype_skip_layers: list[str] = field(default_factory=list) """Layer patterns to skip KV cache quantization. Accepts layer indices (e.g., '0', '2', '4') or attention type names (e.g., 'sliding_window').""" - cpu_kvcache_space_bytes: int | None = None - """(CPU backend only) CPU key-value cache space.""" mamba_page_size_padded: int | None = None """ Optional override for mamba page size; used by hybrid mamba/attention models to ensure exact alignment with attention page size.""" @@ -183,7 +181,6 @@ class CacheConfig: "num_gpu_blocks_override", "enable_prefix_caching", "prefix_caching_hash_algo", - "cpu_kvcache_space_bytes", "mamba_page_size_padded", "user_specified_block_size", "user_specified_mamba_block_size", diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 86f07b1032e..0b41c381bd4 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -6,15 +6,16 @@ import os import platform import subprocess import sys -from dataclasses import dataclass from typing import TYPE_CHECKING -import psutil import torch -from vllm import envs from vllm.logger import init_logger -from vllm.utils.ompmultiprocessing import OMPProcessManager +from vllm.utils.cpu_resource_utils import ( + DEVICE_CONTROL_ENV_VAR, + get_memory_node_info, +) +from vllm.utils.mem_constants import GiB_bytes from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -38,49 +39,13 @@ def get_max_threads(pid=0): raise NotImplementedError("Unsupported OS") -@dataclass -class LogicalCPUInfo: - id: int = -1 - physical_core: int = -1 - numa_node: int = -1 - - @classmethod - def _int(cls, value: str) -> int: - try: - int_value = int(value) - except Exception: - int_value = -1 - return int_value - - @staticmethod - def json_decoder(obj_dict: dict): - id = obj_dict.get("cpu") - physical_core = obj_dict.get("core") - numa_node = obj_dict.get("node") - - if not (id is None or physical_core is None or numa_node is None): - return LogicalCPUInfo( - id=LogicalCPUInfo._int(id), - physical_core=LogicalCPUInfo._int(physical_core), - numa_node=LogicalCPUInfo._int(numa_node), - ) - else: - return obj_dict - - class CpuPlatform(Platform): _enum = PlatformEnum.CPU device_name: str = "cpu" device_type: str = "cpu" dispatch_key: str = "CPU" dist_backend: str = "gloo" - device_control_env_var = "CPU_VISIBLE_MEMORY_NODES" - omp_process_manager = None - # Simultaneous Multithreading (SMT) level for OpenMP: - # 4 on PowerPC, 1 on non-PowerPC architectures - smt = 1 - global_cpu_mask = None - simulate_numa = int(os.environ.get("_SIM_MULTI_NUMA", 0)) + device_control_env_var = DEVICE_CONTROL_ENV_VAR @property def supported_dtypes(self) -> list[torch.dtype]: @@ -123,29 +88,9 @@ class CpuPlatform(Platform): @classmethod def get_device_total_memory(cls, device_id: int = 0) -> int: - from vllm.utils.mem_constants import GiB_bytes - from vllm.utils.mem_utils import format_gib + meminfo = get_memory_node_info(device_id) - kv_cache_space = envs.VLLM_CPU_KVCACHE_SPACE - node_dir = "/sys/devices/system/node" - if kv_cache_space is None: - nodes = ( - [d for d in os.listdir(node_dir) if d.startswith("node")] - if os.path.exists(node_dir) - else [] - ) - num_numa_nodes = len(nodes) or 1 - free_cpu_memory = psutil.virtual_memory().total // num_numa_nodes - DEFAULT_CPU_MEM_UTILIZATION = 0.5 - kv_cache_space = int(free_cpu_memory * DEFAULT_CPU_MEM_UTILIZATION) - logger.warning_once( - "VLLM_CPU_KVCACHE_SPACE not set. Using %s GiB for KV cache.", - format_gib(kv_cache_space), - ) - else: - kv_cache_space *= GiB_bytes - - return kv_cache_space + return meminfo.total_memory @classmethod def set_device(cls, device: torch.device) -> None: @@ -180,6 +125,12 @@ class CpuPlatform(Platform): "otherwise the performance is not optimized." ) + # Lagecy setting + env_key = "VLLM_CPU_KVCACHE_SPACE" + if env_key in os.environ and os.environ[env_key] != "": + kv_cache_space = int(os.environ[env_key]) + cache_config.kv_cache_memory_bytes = kv_cache_space * GiB_bytes + scheduler_config = vllm_config.scheduler_config # async scheduling is not required on CPU scheduler_config.async_scheduling = False @@ -198,8 +149,6 @@ class CpuPlatform(Platform): ) cache_config.cache_dtype = "auto" - cache_config.cpu_kvcache_space_bytes = CpuPlatform.get_device_total_memory() - parallel_config = vllm_config.parallel_config # OMP requires the MP executor to function correctly, UniProc is not # supported as it is not possible to set the OMP environment correctly @@ -278,21 +227,45 @@ class CpuPlatform(Platform): os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1" ld_preload_str = os.getenv("LD_PRELOAD", "") - - # Intel and CLANG OpenMP setting - if "libiomp5.so" in ld_preload_str or "libomp5" in ld_preload_str: - # The time(milliseconds) that a thread should wait after - # completing the execution of a parallel region, before sleeping. - os.environ["KMP_BLOCKTIME"] = "1" - # Prevents the CPU to run into low performance state - os.environ["KMP_TPAUSE"] = "0" - # Provides fine granularity parallelism - os.environ["KMP_FORKJOIN_BARRIER_PATTERN"] = "dist,dist" - os.environ["KMP_PLAIN_BARRIER_PATTERN"] = "dist,dist" - os.environ["KMP_REDUCTION_BARRIER_PATTERN"] = "dist,dist" - cpu_architecture = Platform.get_cpu_architecture() + if ( + platform.system() == "Linux" + and cpu_architecture + in (CpuArchEnum.ARM, CpuArchEnum.POWERPC, CpuArchEnum.X86) + and not ( + "libomp" in ld_preload_str + or "libgomp" in ld_preload_str + or "libiomp" in ld_preload_str + ) + ): + # We need to LD_PRELOAD PyTorch's libgomp, otherwise only + # one core will be properly utilized when we thread-bind + # See: https://github.com/vllm-project/vllm/issues/27369 + # TODO: Remove once: + # https://github.com/pytorch/pytorch/issues/166087 is fixed + + # We need to find the location of PyTorch's libgomp + torch_pkg = os.path.dirname(torch.__file__) + site_root = os.path.dirname(torch_pkg) + # Search both torch.libs and torch/lib - See: + # https://github.com/vllm-project/vllm/issues/30470 + torch_libs_paths = [ + os.path.join(site_root, "torch.libs"), + os.path.join(torch_pkg, "lib"), + ] + pytorch_libgomp_so_candidates = [] + for torch_libs in torch_libs_paths: + pytorch_libgomp_so_candidates.extend( + glob.glob(os.path.join(torch_libs, "libgomp*.so*")) + ) + if pytorch_libgomp_so_candidates: + pytorch_libgomp_so = pytorch_libgomp_so_candidates[0] + if ld_preload_str: + ld_preload_str += ":" + ld_preload_str += pytorch_libgomp_so + os.environ["LD_PRELOAD"] = ld_preload_str + # LD_PRELOAD libtcmalloc, bundled under vllm/libs to reduce # memory allocation overhead if ( @@ -331,13 +304,6 @@ class CpuPlatform(Platform): vllm_config.model_config.max_model_len, vllm_config.scheduler_config.DEFAULT_MAX_NUM_BATCHED_TOKENS, ) - # CI specific "quick" NUMA simulation - split all available CPUs - # into a fake NUMA topology - if os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", None) is not None: - os.environ["_SIM_MULTI_NUMA"] = str( - vllm_config.parallel_config.world_size - * vllm_config.parallel_config._api_process_count - ) @classmethod def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: @@ -345,78 +311,6 @@ class CpuPlatform(Platform): # Move that logic here so block_size is chosen by the backend. pass - @classmethod - def get_omp_manager(cls) -> OMPProcessManager: - # initialise the OMP resource management if need be and return the manager - if cls.omp_process_manager is None: - if cls.get_cpu_architecture() == CpuArchEnum.POWERPC: - cls.smt = 4 - cls.omp_process_manager = OMPProcessManager( - affinity=cls.get_global_cpu_mask(), smt=cls.smt - ) - # we need to fix up the topology returned by the OMP Manager for - # simulated NUMA environments in CI - if cls.simulate_numa > 0: - logger.info( - "Adjusting numa topology to resemble at least %d nodes", - int(cls.simulate_numa), - ) - om = cls.omp_process_manager - while len(om.omp_places) < cls.simulate_numa: - new_omp_places = [] - touched = False - for omp_place in om.omp_places: - if len(omp_place["mask"]) > 1: - touched = True - cpu_list = sorted(list(omp_place["mask"])) - new_omp_places.append( - { - "mask": set(cpu_list[0 : int(len(cpu_list) / 2)]), - "available": True, - } - ) - new_omp_places.append( - { - "mask": set(cpu_list[int(len(cpu_list) / 2) :]), - "available": True, - } - ) - if touched: - om.omp_places = new_omp_places - else: - raise ValueError( - "Cannot split the existing NUMA topology to match " - "simulation requirements" - ) - - return cls.omp_process_manager - - @classmethod - def get_global_cpu_mask(cls) -> set[int]: - # get global cpu mask - if cls.global_cpu_mask is None: - if hasattr(os, "sched_getaffinity"): - cls.global_cpu_mask = os.sched_getaffinity(0) - else: - # macOS does not support sched_getaffinity - cpu_count = os.cpu_count() or 1 - cls.global_cpu_mask = set(range(cpu_count)) - return cls.global_cpu_mask - - @classmethod - def reserve_cpus(cls, reserve: set[int]) -> bool: - # remove CPUs from global mask, for now there is no "release" mechanism - if cls.omp_process_manager is not None: - for place in cls.omp_process_manager.omp_places: - if not place["available"]: - return False - cls.global_cpu_mask = cls.get_global_cpu_mask() - reserve - # reinitialize OMP resource management - cls.omp_process_manager = OMPProcessManager( - affinity=cls.global_cpu_mask, smt=cls.smt - ) - return True - @classmethod def discover_numa_topology(cls) -> list[list[int]]: """ diff --git a/vllm/utils/cpu_resource_utils.py b/vllm/utils/cpu_resource_utils.py new file mode 100644 index 00000000000..bf3d84fb68d --- /dev/null +++ b/vllm/utils/cpu_resource_utils.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import os +import platform +import subprocess +from dataclasses import dataclass +from functools import cache + +import psutil +import regex as re + +DEVICE_CONTROL_ENV_VAR = "CPU_VISIBLE_MEMORY_NODES" + + +@dataclass +class LogicalCPUInfo: + id: int = -1 + physical_core: int = -1 + numa_node: int = -1 + + @classmethod + def _int(cls, value: str) -> int: + try: + int_value = int(value) + except Exception: + int_value = -1 + return int_value + + @staticmethod + def json_decoder(obj_dict: dict): + id = obj_dict.get("cpu") + physical_core = obj_dict.get("core") + numa_node = obj_dict.get("node") + + if not (id is None or physical_core is None or numa_node is None): + return LogicalCPUInfo( + id=LogicalCPUInfo._int(id), + physical_core=LogicalCPUInfo._int(physical_core), + numa_node=LogicalCPUInfo._int(numa_node), + ) + else: + return obj_dict + + +@dataclass +class MemoryNodeInfo: + total_memory: int = -1 + available_memory: int = -1 + + +def get_memory_affinity(pid: int = 0) -> list[int]: + pid = os.getpid() if pid == 0 else pid + path = f"/proc/{pid}/status" + with open(path) as f: + for line in f: + if line.startswith("Mems_allowed_list:"): + # Extract the string part (e.g., "0-1,3") + raw_list = line.split(":")[1].strip() + return parse_id_list(raw_list) + return [] + + +def parse_id_list(raw_str: str) -> list[int]: + """Parses strings like '0-2,4,7-8' into [0, 1, 2, 4, 7, 8]""" + result: list[int] = [] + if not raw_str: + return result + + for part in raw_str.split(","): + if "-" in part: + start, end = map(int, part.split("-")) + result.extend(range(start, end + 1)) + else: + result.append(int(part)) + return sorted(list(set(result))) + + +def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo: + if platform.system() == "Darwin": + # MacOS has no memory node + return MemoryNodeInfo( + total_memory=psutil.virtual_memory().total, + available_memory=psutil.virtual_memory().available, + ) + + meminfo_path = f"/sys/devices/system/node/node{node_id}/meminfo" + if not os.path.exists(meminfo_path): + raise RuntimeError(f"{meminfo_path} doesn't exit.") + + meminfo = {} + with open(meminfo_path) as f: + for line in f: + # Each line looks like: "Node 0 MemTotal: 97421888 kB" + parts = line.split() + key = parts[2].rstrip(":") + # convert to Bytes + value = int(parts[3]) * 1024 + meminfo[key] = value + + total_memory = meminfo["MemTotal"] + free_memory = meminfo["MemFree"] + active_file_memory = meminfo["Active(file)"] + inactive_file_memory = meminfo["Inactive(file)"] + reclaimable_memory = meminfo["SReclaimable"] + available_memory = ( + free_memory + active_file_memory + inactive_file_memory + reclaimable_memory + ) + + return MemoryNodeInfo( + total_memory=total_memory, + available_memory=available_memory, + ) + + +def get_allowed_cpu_list() -> list[LogicalCPUInfo]: + cpu_list = _get_cpu_list() + if platform.system() == "Darwin": + return cpu_list + + global_allowed_cpu_id_list = os.sched_getaffinity(0) + logical_cpu_list = [x for x in cpu_list if x.id in global_allowed_cpu_id_list] + + return logical_cpu_list + + +def get_visible_memory_node() -> list[int]: + if platform.system() == "Darwin": + return [0] + + allowed_memory_node_list = get_memory_affinity() + + env_key = DEVICE_CONTROL_ENV_VAR + if ( + ("VLLM_CPU_SIM_MULTI_NUMA" not in os.environ) + and env_key in os.environ + and os.environ[env_key] != "" + ): + visible_nodes = [int(s) for s in os.environ[env_key].split(",")] + visible_nodes = [ + node for node in visible_nodes if node in allowed_memory_node_list + ] + return visible_nodes + + return allowed_memory_node_list + + +@cache +def _get_cpu_list() -> list[LogicalCPUInfo]: + if platform.system() == "Darwin": + # For MacOS, no user-level CPU affinity and SMT, return all CPUs + cpu_count = os.cpu_count() + assert cpu_count + return [LogicalCPUInfo(i, i, 0) for i in range(cpu_count)] + + lscpu_output = subprocess.check_output( + "lscpu -J -e=CPU,CORE,NODE", shell=True, text=True + ) + + # For platform without NUMA, replace '-' to '0' + lscpu_output = re.sub(r'"node":\s*-\s*(,|\n)', r'"node": 0\1', lscpu_output) + + logical_cpu_list: list[LogicalCPUInfo] = json.loads( + lscpu_output, object_hook=LogicalCPUInfo.json_decoder + )["cpus"] + + # Filter CPUs with invalid attributes + logical_cpu_list = [ + x for x in logical_cpu_list if -1 not in (x.id, x.physical_core, x.numa_node) + ] + + return logical_cpu_list diff --git a/vllm/utils/ompmultiprocessing.py b/vllm/utils/ompmultiprocessing.py index d601e39e15f..c3c607ea90b 100644 --- a/vllm/utils/ompmultiprocessing.py +++ b/vllm/utils/ompmultiprocessing.py @@ -5,196 +5,280 @@ Copyright (c) 2026 Red Hat Inc Copyright (c) 2026 Cambridge Greys Ltd """ -import json import os -import platform -import subprocess +from collections.abc import Callable +from contextlib import contextmanager +from typing import TYPE_CHECKING + +import vllm.utils.cpu_resource_utils as cr_utils +from vllm import envs +from vllm.logger import init_logger +from vllm.platforms import CpuArchEnum, current_platform +from vllm.utils.cpu_resource_utils import LogicalCPUInfo + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) -def _int(arg): - """Relaxed parsing of ints which handles a - instead of a number. - The lscpu json may contain that for nodes in some cases. If that - is the case we parse it to zero - """ - try: - if int(arg) >= 0: - return int(arg) - except ValueError: - pass - return 0 - - -def parse_mask(mask): - """Expand a X-Y,Z list""" - result = [] - for token in mask.split(","): - try: - start, finish = token.split("-") - if int(start) > int(finish): - raise IndexError("Invalid Indexes for cpu ranges") - for cpu in range(int(start), int(finish) + 1): - result.append(cpu) - except ValueError: - result.append(int(token)) - return set(result) - - -def _get_default_affinity() -> set[int]: - """Get the set of CPUs the process is allowed to run on.""" - if hasattr(os, "sched_getaffinity"): - return os.sched_getaffinity(0) - # macOS does not support sched_getaffinity; fall back to cpu_count - cpu_count = os.cpu_count() or 1 - return set(range(cpu_count)) - - -def _get_cpu_topology_json() -> bytes: - """Get CPU topology as JSON. - - On Linux this uses ``lscpu -Je``. On other platforms (e.g. macOS) we - synthesize a simple topology where every logical CPU is its own core - on NUMA node 0, which is sufficient for the OMP place-list builder. - """ - if platform.system() == "Linux": - return subprocess.run(["lscpu", "-Je"], check=True, capture_output=True).stdout - - # Fallback for non-Linux (macOS, etc.) - cpu_count = os.cpu_count() or 1 - cpus = [] - for i in range(cpu_count): - cpus.append({"cpu": str(i), "core": str(i), "node": "0"}) - return json.dumps({"cpus": cpus}).encode() - - -def enumerate_resources(resource_map, mask=None, allowed=None): - """Enumerate system resources""" - if allowed is None: - allowed = _get_default_affinity() - if mask is not None: - allowed = allowed & mask - - try: - allowed_nodes = parse_mask(os.environ["CPU_VISIBLE_MEMORY_NODES"]) - except KeyError: - allowed_nodes = None - - lscpu: dict[str, dict] = {"cpus": {}, "cores": {}, "nodes": {}} - for cpu in resource_map["cpus"]: - cpunum = int(cpu["cpu"]) - if ( - cpunum in allowed - and cpunum >= 0 - and (allowed_nodes is None or _int(cpu["node"]) in allowed_nodes) - ): - lscpu["cpus"][cpunum] = [cpu] - core = _int(cpu["core"]) - if lscpu["cores"].get(core, None) is None: - lscpu["cores"][core] = [cpu] - else: - lscpu["cores"][core].append(cpu) - node = _int(cpu["node"]) - if lscpu["nodes"].get(node, None) is None: - lscpu["nodes"][node] = [cpu] - else: - lscpu["nodes"][node].append(cpu) - return lscpu - - -def produce_cpu_list(cpus, smt=1): - """Produce a CPU list with/without SMT pairs - main cpu list case""" - mask: list[int] = [] - for key, value in cpus.items(): - exists = 0 - for cpu in mask: - if cpu == value[0]["core"]: - exists += 1 - break - if exists < smt: - mask.append(int(key)) - return {"mask": set(mask), "available": True} - - -def produce_cpu_sublist(scpus, smt=1): - """Produce a CPU list with/without SMT pairs - resource leaf case""" - cpu_list: list[dict] = [] - for value in scpus: - exists = 0 - for cpu in cpu_list: - if int(cpu["core"]) == int(value["core"]): - exists += 1 - break - if exists < smt: - cpu_list.append(value) - mask = [] - for cpu in cpu_list: - mask.append(int(cpu["cpu"])) - - return {"mask": set(mask), "available": True} - - -def create_omp_places(resources, strategy, smt=True): - """Parse CPU topology and generate possible CPU masks""" - omp_places = [] - if strategy == "all": - omp_places.append(produce_cpu_list(resources["cpus"], smt)) - elif strategy == "cores": - for value in resources["cores"].values(): - omp_places.append(produce_cpu_sublist(value, smt)) - elif strategy == "nodes": - for value in resources["nodes"].values(): - omp_places.append(produce_cpu_sublist(value, smt)) - else: - raise NotImplementedError("Unknown strategy") - - return omp_places - - -# pylint: disable=too-few-public-methods class OMPProcessManager: - """OMP aware wrapper to run mp Process()""" + def __init__(self, config: "VllmConfig"): + if not current_platform.is_cpu(): + return - def __init__(self, strategy="nodes", smt=1, mock=None, affinity=None): - self.strategy = strategy - self.smt = smt - self.omp_places = [] - vllm_mask = os.environ.get("VLLM_CPU_OMP_THREADS_BIND", None) - self.setup_omp = vllm_mask != "nobind" - if self.setup_omp: - omp_places = [] - if vllm_mask is not None: - masks = [] - for spec in vllm_mask.split("|"): - masks.append(parse_mask(spec)) + self.local_world_size = config.parallel_config.local_world_size + self.local_dp_rank = config.parallel_config.data_parallel_rank_local + # This is a bit tricky because the internal DP size + # is always 1 for non-MoE models + self.internal_dp_size = config.parallel_config._api_process_count + + self.simulate_multi_node = os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", "0") != "0" + ld_preload_str = os.getenv("LD_PRELOAD", "") + self.use_iomp = "libiomp" in ld_preload_str or "libomp" in ld_preload_str + self.use_gomp = "libgomp" in ld_preload_str + + assert not (self.use_iomp and self.use_gomp) + + # at least reserve 1/local_world_size(for ARM) core for scheduler + # proc as always use MP executor + # TODO: make scheduler proc sleep when idle + self.reserve_cpu_num = ( + self.local_world_size + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM + else 1 + ) + # reserve at one more core for nixl_connector under p/d case + if config.kv_transfer_config: + self.reserve_cpu_num += 1 + + if envs.VLLM_CPU_NUM_OF_RESERVED_CPU is not None: + if self.reserve_cpu_num > envs.VLLM_CPU_NUM_OF_RESERVED_CPU: + msg = ( + f"VLLM_CPU_NUM_OF_RESERVED_CPU is less than " + "the minimum requirement" + f": {self.reserve_cpu_num} cores" + ) + logger.warning(msg=msg) + self.reserve_cpu_num = envs.VLLM_CPU_NUM_OF_RESERVED_CPU + + self._parse_omp_threads_bind_env() + + assert not self.simulate_multi_node or self.auto_setup + + @contextmanager + def configure_omp_envs(self, rank: int, local_rank: int): + if not current_platform.is_cpu() or self.skip_setup: + yield + return + + envs_dict = {} + cpu_list = [str(i) for i in self.cpu_lists[local_rank]] + envs_dict["OMP_NUM_THREADS"] = str(len(cpu_list)) + if self.use_iomp: + # set IOMP envs + cpu_list_str = ",".join(cpu_list) + envs_dict["KMP_AFFINITY"] = ( + f"granularity=fine,explicit,proclist=[{cpu_list_str}]" + ) + # The time(milliseconds) that a thread should wait after + # completing the execution of a parallel region, before sleeping. + envs_dict["KMP_BLOCKTIME"] = "1" + # Prevents the CPU to run into low performance state + envs_dict["KMP_TPAUSE"] = "0" + # Provides fine granularity parallelism + envs_dict["KMP_FORKJOIN_BARRIER_PATTERN"] = "dist,dist" + envs_dict["KMP_PLAIN_BARRIER_PATTERN"] = "dist,dist" + envs_dict["KMP_REDUCTION_BARRIER_PATTERN"] = "dist,dist" + elif self.use_gomp: + # set GOMP envs + # likes '0 1 2 ...' + cpu_list_str = " ".join(cpu_list) + envs_dict["GOMP_CPU_AFFINITY"] = cpu_list_str + else: + # set OMP envs + # likes '{0,1,2,...}' + cpu_list_str = ",".join(cpu_list) + envs_dict["OMP_PLACES"] = f"{{{cpu_list_str}}}" + envs_dict["OMP_PROC_BIND"] = "true" + + # backup envs + old_envs_dict = {} + for k in envs_dict: + old_envs_dict[k] = os.environ.get(k) + + try: + # set envs + for k, v in envs_dict.items(): + os.environ[k] = v + yield + finally: + # restore old envs + for k, v in old_envs_dict.items(): # type: ignore + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + + def _parse_omp_threads_bind_env(self): + vllm_mask = envs.VLLM_CPU_OMP_THREADS_BIND + self.skip_setup = vllm_mask == "nobind" + self.auto_setup = vllm_mask == "auto" + self.reserved_cpu_list = [] + self.cpu_lists = [] + + if self.auto_setup: + # auto generate CPU lists + cpu_arch = current_platform.get_cpu_architecture() + if cpu_arch == CpuArchEnum.POWERPC: + # For POWERPC SMT-8/4/2 + cpu_list, reserve_list = self._get_autobind_cpu_ids( + lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4] + ) + elif cpu_arch in (CpuArchEnum.X86, CpuArchEnum.S390X): + # For x86/S390X SMT-2, use 1 logical CPU per physical core + cpu_list, reserve_list = self._get_autobind_cpu_ids( + lambda cpus: cpus[-1:] + ) + elif cpu_arch == CpuArchEnum.ARM: + # For AArch64, no SMT, use all logical CPU + cpu_list, reserve_list = self._get_autobind_cpu_ids(lambda cpus: cpus) else: - masks = [None] - if mock is None: - data = _get_cpu_topology_json() - else: - with open(mock, mode="rb") as jf: - data = jf.read() - lscpu = json.loads(data) - for mask in masks: - resources = enumerate_resources(lscpu, mask, affinity) - omp_places.extend(create_omp_places(resources, strategy, smt)) - self.omp_places = sorted( - omp_places, - key=lambda p: "{:04d}-{:04d}".format(len(p["mask"]), max(p["mask"])), - reverse=True, + cpu_list, reserve_list = [], [] + raise RuntimeError(f"{cpu_arch} doesn't support auto CPU binding.") + + for item in cpu_list: + self.cpu_lists.append([x.id for x in item]) + self.reserved_cpu_list = [x.id for x in reserve_list] + elif not self.skip_setup: + # user defined CPU lists + omp_cpuids_list = vllm_mask.split("|") + if self.local_dp_rank is not None: + local_dp_rank = self.local_dp_rank + world_size = self.local_world_size + # Rank mapping [DP, PP, TP] + omp_cpuids_list = omp_cpuids_list[ + local_dp_rank * world_size : (local_dp_rank + 1) * world_size + ] + + assert len(omp_cpuids_list) == self.local_world_size, ( + "Given " + f"number of CPU id list {omp_cpuids_list} doesn't match " + f"local world size {self.local_world_size}." ) - def run(self, what, *args, **kwargs): - """Run arg with correct OMP environment""" - if self.setup_omp: - for place in self.omp_places: - if place["available"]: - reserve = int(os.environ.get("VLLM_CPU_NUM_OF_RESERVED_CPU", 0)) - place["available"] = False - # pylint: disable=consider-using-f-string - os.environ["OMP_PLACES"] = "{}".format(place["mask"]) - os.environ["OMP_NUM_THREADS"] = "{}".format( - len(place["mask"]) - reserve - ) - os.environ["OMP_PROC_BIND"] = "TRUE" - return what(*args, **kwargs) - raise IndexError("Out of OMP places") - return what(*args, **kwargs) + # parse CPU list strings like "5,2-4" to [5, 2, 3, 4] + self.cpu_lists = [cr_utils.parse_id_list(s) for s in omp_cpuids_list] + else: + # skip + self.cpu_lists = [] + + msg = "OpenMP thread binding info: \n" + for i in range(self.local_world_size): + msg += f"\tlocal_rank={i}, core ids={self.cpu_lists[i]}\n" + msg += f"\treserved_cpus={self.reserved_cpu_list}" + logger.info(msg) + + def _get_autobind_cpu_ids( + self, cpu_selector: Callable[[list[LogicalCPUInfo]], list[LogicalCPUInfo]] + ) -> tuple[list[list[LogicalCPUInfo]], list[LogicalCPUInfo]]: + """ + Return CPU ids to bind based on NUMA nodes, and CPU ids reserved for + other processes. + Currently for rank N, only CPU ids on the N-th node in available NUMA + node list will be selected. + Args: + cpu_selector: a callable object to select CPUs from a CPU list + of a physical core. The input is a LogicalCPUInfo list contains + logical CPUs of a physical CPU, sorted by the LogicalCPUInfo.id. + A selected LogicalCPUInfo list should be returned. + """ + + # this memory node list has been sliced for DP offset + allowed_numa_nodes = cr_utils.get_visible_memory_node() + logical_cpu_list = cr_utils.get_allowed_cpu_list() + + local_world_size = self.local_world_size + assert ( + len(allowed_numa_nodes) >= local_world_size or self.simulate_multi_node + ), ( + f"Not enough allowed NUMA nodes to bind threads of " + f"{local_world_size} local CPUWorkers. " + f"Allowed NUMA nodes are {allowed_numa_nodes}. " + "Please try to bind threads manually or decrease DP/TP/PP." + ) + + # Generate OMP CPU list for each rank + cpu_lists_of_ranks = [] + reserved_cpu_list = [] + total_cpu_num = 0 + for local_rank in range(self.local_world_size): + if not self.simulate_multi_node: + selected_numa_node = allowed_numa_nodes[local_rank] + selected_logical_cpu_list = [ + x for x in logical_cpu_list if x.numa_node == selected_numa_node + ] + else: + world_size_across_dp = self.local_world_size * self.internal_dp_size + assert len(logical_cpu_list) >= world_size_across_dp + selected_logical_cpu_list = sorted( + logical_cpu_list, key=lambda x: x.numa_node + ) + sim_cpu_num_per_node = ( + len(selected_logical_cpu_list) // world_size_across_dp + ) + assert self.local_dp_rank is not None + start_idx = ( + local_rank + self.local_world_size * self.local_dp_rank + ) * sim_cpu_num_per_node + selected_logical_cpu_list = selected_logical_cpu_list[ + start_idx : (start_idx + sim_cpu_num_per_node) + ] + + # Select logical CPUs on same physical cores via cpu_selector + core_to_cpus: dict[int, list[LogicalCPUInfo]] = {} + for cpu_info in selected_logical_cpu_list: + if cpu_info.physical_core not in core_to_cpus: + core_to_cpus[cpu_info.physical_core] = [] + core_to_cpus[cpu_info.physical_core].append(cpu_info) + selected_logical_cpu_list = [] + for cpu_list in core_to_cpus.values(): + cpu_list = sorted(cpu_list, key=lambda x: x.id) + selected_logical_cpu_list.extend(cpu_selector(cpu_list)) + + # sort selected cores based on core id + selected_logical_cpu_list = sorted( + selected_logical_cpu_list, key=lambda x: x.id + ) + + cpu_lists_of_ranks.append(selected_logical_cpu_list) + total_cpu_num += len(selected_logical_cpu_list) + + # Reserve CPUs for other processes + if total_cpu_num <= self.reserve_cpu_num: + logger.warning( + "Selected CPU core number (%s) " + "should be greater than reserved CPU core " + "number (%s).", + total_cpu_num, + self.reserve_cpu_num, + ) + return cpu_lists_of_ranks, [] + + reserve_num_per_rank = [ + self.reserve_cpu_num // self.local_world_size + ] * self.local_world_size + # last rank first + for i in range( + self.local_world_size - 1, + self.local_world_size - 1 - self.reserve_cpu_num % self.local_world_size, + -1, + ): + reserve_num_per_rank[i] += 1 + for i in range(self.local_world_size): + num = reserve_num_per_rank[i] + if num > 0: + reserved_cpu_list.extend(cpu_lists_of_ranks[i][-num:]) + cpu_lists_of_ranks[i] = cpu_lists_of_ranks[i][:-num] + + return cpu_lists_of_ranks, reserved_cpu_list diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 3eccce722b5..52969783f09 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -51,6 +51,7 @@ from vllm.utils.network_utils import ( get_loopback_ip, get_open_port, ) +from vllm.utils.ompmultiprocessing import OMPProcessManager from vllm.utils.system_utils import ( _maybe_force_spawn, decorate_logs, @@ -169,24 +170,14 @@ class MultiprocExecutor(Executor): [] if context.get_start_method() == "fork" else None ) + # For CPU backend only, to setup OpenMP threads affinity + cpu_omp_manager = OMPProcessManager(self.vllm_config) for local_rank in range(self.local_world_size): global_rank = global_start_rank + local_rank is_driver_worker = self._is_driver_worker(global_rank) - if current_platform.is_cpu(): - om = current_platform.get_omp_manager() - logger.info("Configured OMP PLACES %s", str(om.omp_places)) - unready_worker_handle = om.run( - WorkerProc.make_worker_process, - vllm_config=self.vllm_config, - local_rank=local_rank, - rank=global_rank, - distributed_init_method=distributed_init_method, - input_shm_handle=scheduler_output_handle, - shared_worker_lock=shared_worker_lock, - is_driver_worker=is_driver_worker, - inherited_fds=inherited_fds, - ) - else: + with cpu_omp_manager.configure_omp_envs( + rank=global_rank, local_rank=local_rank + ): unready_worker_handle = WorkerProc.make_worker_process( vllm_config=self.vllm_config, local_rank=local_rank, diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 7067ef62a75..d7c52d58a5e 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/vllm/v1/worker/cpu_model_runner.py @@ -116,21 +116,7 @@ class CPUModelRunner(GPUModelRunner): logger.info("Warming up model for the compilation...") # Only generate graph for the generic shape with _set_global_compilation_settings(self.vllm_config): - self._dummy_run( - min( - max(16, self.max_num_reqs), - self.scheduler_config.max_num_batched_tokens, - ) - ) - - # Warm up drafter for speculative decoding - if self.speculative_config and (self.speculative_config.uses_draft_model()): - from vllm.v1.spec_decode.draft_model import DraftModelProposer - - if isinstance(self.drafter, (DraftModelProposer)): - logger.info("Warming up drafter model...") - self.drafter.dummy_run(max(16, self.max_num_reqs)) - + self.profile_run() logger.info("Warming up done.") def initialize_kv_cache( diff --git a/vllm/v1/worker/cpu_worker.py b/vllm/v1/worker/cpu_worker.py index 1d85a94a4e1..ce144c3aa21 100644 --- a/vllm/v1/worker/cpu_worker.py +++ b/vllm/v1/worker/cpu_worker.py @@ -1,15 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math import os import sys from typing import Any +import psutil import torch from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.platforms import CpuArchEnum, current_platform from vllm.profiler.wrapper import TorchProfilerWrapper +from vllm.utils.cpu_resource_utils import ( + get_allowed_cpu_list, + get_memory_node_info, + get_visible_memory_node, +) +from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import set_random_seed from vllm.v1.worker.cpu_model_runner import CPUModelRunner from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment @@ -27,6 +35,46 @@ class CPUWorker(Worker): distributed_init_method: str, is_driver_worker: bool = False, ): + # TODO: use numactl for process setup + # TODO: optimize for `interleaved` policy + # Bind memory node + allowed_memory_nodes = get_visible_memory_node() + allowed_cpu_list = get_allowed_cpu_list() + cpu_core = allowed_cpu_list[0] + + # TODO: some CI hosts are not correctly set, change to assertion + # after fix + if cpu_core.numa_node not in allowed_memory_nodes: + logger.warning( + "Node %s is not in available memory nodes %s.", + cpu_core.numa_node, + allowed_memory_nodes, + ) + + torch.ops._C.init_cpu_memory_env([cpu_core.numa_node]) + + memory_status = get_memory_node_info(cpu_core.numa_node) + memory_fraction = vllm_config.cache_config.gpu_memory_utilization + self.requested_cpu_memory = math.ceil( + memory_status.total_memory * memory_fraction + ) + available_memory = memory_status.available_memory + + if ( + vllm_config.cache_config.kv_cache_memory_bytes is None + and self.requested_cpu_memory > available_memory + ): + raise ValueError( + f"Available memory on node {cpu_core.numa_node} " + f"({format_gib(available_memory)}/" + f"{format_gib(memory_status.total_memory)} GiB) on startup " + f"is less than desired CPU memory utilization " + f"({vllm_config.cache_config.gpu_memory_utilization}, " + f"{format_gib(self.requested_cpu_memory)} GiB). " + "Decrease --gpu-memory-utilization" + f" or reduce CPU memory used by other processes." + ) + super().__init__( vllm_config, local_rank, @@ -103,13 +151,69 @@ class CPUWorker(Worker): pass def determine_available_memory(self) -> int: - return self.cache_config.cpu_kvcache_space_bytes or 0 + self.model_runner.warming_up_model() + + allowed_cpu_list = get_allowed_cpu_list() + cpu_core = allowed_cpu_list[0] + + memory_status = get_memory_node_info(cpu_core.numa_node) + available_memory = memory_status.available_memory + explicit_kv_cache_size = self.cache_config.kv_cache_memory_bytes + + kv_cache_size = None + msg = None + if explicit_kv_cache_size is not None: + if explicit_kv_cache_size > available_memory: + raise ValueError( + f"Available memory on node {cpu_core.numa_node} " + f"({format_gib(available_memory)}/" + f"{format_gib(memory_status.total_memory)} GiB) on kv cache" + f" allocation is less than requested memory for kv " + f"({format_gib(explicit_kv_cache_size)} GiB). " + "Decrease --kv-cache-memory-bytes, VLLM_CPU_KVCACHE_SPACE, " + "or reduce CPU memory used by other processes." + ) + kv_cache_size = explicit_kv_cache_size + msg = ( + f"Explicitly set ({format_gib(kv_cache_size)}/" + f"{format_gib(memory_status.total_memory)}) GiB for KV cache " + f"on node {cpu_core.numa_node}." + ) + else: + consumed_memory = psutil.Process(os.getpid()).memory_info().rss + requested_memory_for_kv = int(self.requested_cpu_memory - consumed_memory) + if ( + requested_memory_for_kv <= 0 + or requested_memory_for_kv > available_memory + ): + raise ValueError( + f"Available memory on node {cpu_core.numa_node} " + f"({format_gib(available_memory)}/" + f"{format_gib(memory_status.total_memory)} GiB) on kv cache" + f" allocation is less than requested memory for kv " + f"({format_gib(requested_memory_for_kv)}/" + f"{format_gib(self.requested_cpu_memory)} GiB). " + "Reduce CPU memory used by other processes." + ) + kv_cache_size = requested_memory_for_kv + msg = ( + f"Auto set ({format_gib(kv_cache_size)}/" + f"{format_gib(memory_status.total_memory)}) GiB for KV cache " + f"on node {cpu_core.numa_node}, with " + f"{format_gib(self.requested_cpu_memory)} GiB requested memory" + f" for the worker. {format_gib(consumed_memory)} GiB" + f" memory was consumed by non-kv usages." + ) + + logger.info(msg) + + return kv_cache_size def compile_or_warm_up_model(self) -> CompilationTimes: # Reset the seed to ensure that the random state is not affected by # the model initialization and profiling. set_random_seed(self.model_config.seed) - self.model_runner.warming_up_model() + # Note: the model has been compiled in determine_available_memory() return CompilationTimes( language_model=self.compilation_config.compilation_time, encoder=self.compilation_config.encoder_compilation_time, From 7a51b3e415eeb2954e471e8b5a3898047276e492 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Fri, 17 Apr 2026 21:34:55 +0800 Subject: [PATCH 084/150] [Bugfix] Fix empty delta detection in Qwen3XMLToolParser streaming (#40090) Signed-off-by: chaunceyjiang --- tests/tool_parsers/test_qwen3xml_tool_parser.py | 3 --- vllm/tool_parsers/qwen3xml_tool_parser.py | 11 +++++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/tool_parsers/test_qwen3xml_tool_parser.py b/tests/tool_parsers/test_qwen3xml_tool_parser.py index 3771b8afd24..1ea9a1d65c0 100644 --- a/tests/tool_parsers/test_qwen3xml_tool_parser.py +++ b/tests/tool_parsers/test_qwen3xml_tool_parser.py @@ -64,9 +64,6 @@ class TestQwen3xmlToolParser(ToolParserTests): "test_empty_arguments": "Qwen3XML streaming has systematic issues", "test_surrounding_text": "Qwen3XML streaming has systematic issues", "test_escaped_strings": "Qwen3XML streaming has systematic issues", - "test_malformed_input": ( - "Qwen3XML parser is lenient with malformed input" - ), "test_streaming_reconstruction": ( "Qwen3XML streaming reconstruction has known issues" ), diff --git a/vllm/tool_parsers/qwen3xml_tool_parser.py b/vllm/tool_parsers/qwen3xml_tool_parser.py index 4ecb9666886..8ee10dcbc9e 100644 --- a/vllm/tool_parsers/qwen3xml_tool_parser.py +++ b/vllm/tool_parsers/qwen3xml_tool_parser.py @@ -1258,11 +1258,11 @@ class Qwen3XMLToolParser(ToolParser): return None # Parse the delta text and get the result - result = self.parser.parse_single_streaming_chunks(delta_text) + delta = self.parser.parse_single_streaming_chunks(delta_text) # Update tool call tracking arrays based on incremental parsing results - if result and result.tool_calls: - for tool_call in result.tool_calls: + if delta and delta.tool_calls: + for tool_call in delta.tool_calls: if tool_call.function: tool_index = ( tool_call.index @@ -1292,4 +1292,7 @@ class Qwen3XMLToolParser(ToolParser): self.streamed_args_for_tool[tool_index] += ( tool_call.function.arguments ) - return result + if delta.content is None and not delta.tool_calls and delta.reasoning is None: + # If no content and no tool calls, return None to indicate no update + return None + return delta From 70770268c3953db13aca974242e04bd4be73491c Mon Sep 17 00:00:00 2001 From: Ben Browning Date: Fri, 17 Apr 2026 09:51:48 -0400 Subject: [PATCH 085/150] Add @bbrowning to CODEOWNERS (#40141) Signed-off-by: Ben Browning --- .github/CODEOWNERS | 12 ++++++------ docs/governance/committers.md | 3 ++- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index e272bf9c477..a20c5e7e9dc 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -44,9 +44,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson /vllm/pooling_params.py @noooop @DarkLight1337 /vllm/tokenizers @DarkLight1337 @njhill /vllm/renderers @DarkLight1337 @njhill -/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 -/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 -/vllm/parser @aarnphm @chaunceyjiang @sfeng33 +/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning # vLLM V1 /vllm/v1/attention @LucasWilkinson @MatthewBonanni @@ -93,9 +93,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson /tests/v1/kv_connector @ApostaC @orozery /tests/v1/kv_offload @ApostaC @orozery /tests/v1/determinism @yewentao256 -/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 -/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 -/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 +/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning # Transformers modeling backend /vllm/model_executor/models/transformers @hmellor diff --git a/docs/governance/committers.md b/docs/governance/committers.md index 42c8a8ca2fe..cff041866b2 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -14,6 +14,7 @@ Sorted alphabetically by GitHub handle: - [@aarnphm](https://github.com/aarnphm): Structured output - [@alexm-redhat](https://github.com/alexm-redhat): Performance - [@ApostaC](https://github.com/ApostaC): Connectors, offloading +- [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser - [@benchislett](https://github.com/benchislett): Engine core and spec decode - [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration - [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser @@ -121,7 +122,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - State space models: The state space models implementation in vLLM - @tdoublep, @tlrmchlsmth - Reasoning and tool calling parsers - - @chaunceyjiang, @aarnphm, @sfeng33 + - @chaunceyjiang, @aarnphm, @sfeng33, @bbrowning ### Entrypoints From 6b2b7bd0ebd43ef756632d2142ce974929f05d8f Mon Sep 17 00:00:00 2001 From: sychen52 <41452870+sychen52@users.noreply.github.com> Date: Fri, 17 Apr 2026 07:28:00 -0700 Subject: [PATCH 086/150] Add nvfp4 support to reshape_and_cache_flash (#37332) Signed-off-by: Shiyang Chen --- CMakeLists.txt | 14 + csrc/cache_kernels.cu | 24 +- csrc/nvfp4_kv_cache_kernels.cu | 275 ++++++++++++++++++ tests/kernels/attention/test_cache.py | 120 ++++++-- tests/kernels/quantization/nvfp4_utils.py | 54 ++++ vllm/config/cache.py | 1 + .../layers/attention/attention.py | 6 +- vllm/utils/torch_utils.py | 134 ++++++++- vllm/v1/attention/backends/flashinfer.py | 72 ++++- vllm/v1/kv_cache_interface.py | 30 +- 10 files changed, 679 insertions(+), 51 deletions(-) create mode 100644 csrc/nvfp4_kv_cache_kernels.cu diff --git a/CMakeLists.txt b/CMakeLists.txt index f24c12eff83..fd4314bec52 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -923,6 +923,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") SRCS "${SRCS}" CUDA_ARCHS "${FP4_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly + # from cache_kernels.cu, so it belongs in _C rather than _C_stable. + set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu") + set_gencode_flags_for_srcs( + SRCS "${NVFP4_KV_SRC}" + CUDA_ARCHS "${FP4_ARCHS}") + target_sources(_C PRIVATE ${NVFP4_KV_SRC}) + target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1) list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1") message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}") @@ -949,6 +957,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") SRCS "${SRCS}" CUDA_ARCHS "${FP4_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu") + set_gencode_flags_for_srcs( + SRCS "${NVFP4_KV_SRC}" + CUDA_ARCHS "${FP4_ARCHS}") + target_sources(_C PRIVATE ${NVFP4_KV_SRC}) + target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1) list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1") message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}") diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index d1cfbaeb05d..6bea5abc3df 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -724,6 +724,28 @@ void reshape_and_cache_flash( int num_tokens = slot_mapping.size(0); int num_heads = key.size(1); int head_size = key.size(2); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + if (kv_cache_dtype == "nvfp4") { +#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120) + // NVFP4 dispatch is compiled separately for SM100+. + extern void reshape_and_cache_nvfp4_dispatch( + torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache, + torch::Tensor & value_cache, torch::Tensor & slot_mapping, + torch::Tensor & k_scale, torch::Tensor & v_scale); + reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache, + slot_mapping, k_scale, v_scale); + return; +#else + TORCH_CHECK(false, + "NVFP4 KV cache requires SM100+ (Blackwell). " + "Please rebuild vllm with a Blackwell-compatible CUDA target."); +#endif + } + + // Original FP8/auto path. int block_size = key_cache.size(1); int64_t key_stride = key.stride(0); @@ -741,8 +763,6 @@ void reshape_and_cache_flash( dim3 grid(num_tokens); dim3 block(std::min(num_heads * head_size, 512)); - const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype, CALL_RESHAPE_AND_CACHE_FLASH); diff --git a/csrc/nvfp4_kv_cache_kernels.cu b/csrc/nvfp4_kv_cache_kernels.cu new file mode 100644 index 00000000000..d6aa715c203 --- /dev/null +++ b/csrc/nvfp4_kv_cache_kernels.cu @@ -0,0 +1,275 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +// NVFP4 KV cache store kernel. +// Quantizes bf16 key/value to packed FP4 + FP8 block scales and writes them +// into the paged KV cache. +// +// Per page layout: [K_data | K_scale | V_data | V_scale] +// Both data and scale regions are contiguous per head, enabling direct +// TMA descriptor use. +// +// Reuses device functions from nvfp4_utils.cuh: +// - cvt_warp_fp16_to_fp4() for bf16 → fp4 quantization + block scale +// - pack_fp4() for packing float pairs to fp4 +// - reciprocal_approximate_ftz() for fast reciprocal + +#define NVFP4_ENABLE_ELTS16 1 +#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh" + +#include +#include +#include + +#include "dispatch_utils.h" + +namespace vllm { + +// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel. +// The swizzle pattern for HND layout is: +// [T//4, 4, 4, S//4] → permute(0, 2, 3, 1) → reshape to [T, S] +// where T = block_size (page_size), S = scale_dim = head_size // 16. +// +// For a linear (t, s) position, the swizzled position is: +// swizzled_t = (t / 4) * 4 + (s / (S / 4)) +// swizzled_s = (s % (S / 4)) * 4 + (t % 4) +__device__ __forceinline__ int swizzle_scale_offset(int t, int s, + int scale_dim) { + int s_group = scale_dim / 4; + int swizzled_t = (t / 4) * 4 + (s / s_group); + int swizzled_s = (s % s_group) * 4 + (t % 4); + return swizzled_t * scale_dim + swizzled_s; +} + +// Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache. +// +// Takes separate data and scale cache pointers for K and V. +// Within each KV side, data and scale are separate contiguous regions. +// +// Threading: one CUDA block per token, threads process heads and +// groups of 16 elements within each head. +template +__global__ void reshape_and_cache_nvfp4_kernel( + const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size] + const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size] + uint8_t* __restrict__ key_data_cache, // data region for K + uint8_t* __restrict__ value_data_cache, // data region for V + uint8_t* __restrict__ key_scale_cache, // scale region for K + uint8_t* __restrict__ value_scale_cache, // scale region for V + const int64_t* __restrict__ slot_mapping, // [num_actual_tokens] + const float* __restrict__ k_scale_ptr, // pointer to checkpoint k_scale + const float* __restrict__ v_scale_ptr, // pointer to checkpoint v_scale + const int64_t key_stride, // key.stride(0) in elements + const int64_t value_stride, // value.stride(0) in elements + const int num_heads, const int head_size, const int block_size, + const int64_t data_block_stride, // data cache stride for dim 0 + const int64_t data_head_stride, // data cache stride for heads + const int64_t data_block_offset_stride, // data cache stride for tokens + const int64_t scale_block_stride, // scale cache stride for dim 0 + const int64_t scale_head_stride, // scale cache stride for heads + const int64_t scale_block_offset_stride // scale cache stride for tokens +) { + using CudaType = typename CUDATypeConverter::Type; + using PVec = PackedVec; + + static constexpr int ELTS = CVT_FP4_ELTS_PER_THREAD; // 16 or 8 + static constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / ELTS; + + const int64_t token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[token_idx]; + if (slot_idx < 0) return; + + const int64_t block_idx = slot_idx / block_size; + const int block_offset = static_cast(slot_idx % block_size); + + const int scale_dim = head_size / 16; + const int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE; + + const int total_groups = num_heads * groups_per_head; + const int tid = threadIdx.x; + const int num_thread_groups = blockDim.x / THREADS_PER_SF; + const int tg_id = tid / THREADS_PER_SF; + const int tg_lane = tid % THREADS_PER_SF; + + // Process both K (kv=0) and V (kv=1) +#pragma unroll + for (int kv = 0; kv < 2; kv++) { + const scalar_t* __restrict__ src = (kv == 0) ? key : value; + const float global_scale = 1.0f / ((kv == 0) ? *k_scale_ptr : *v_scale_ptr); + const int64_t src_stride = (kv == 0) ? key_stride : value_stride; + uint8_t* __restrict__ data_cache = + (kv == 0) ? key_data_cache : value_data_cache; + uint8_t* __restrict__ sc_cache = + (kv == 0) ? key_scale_cache : value_scale_cache; + + // Source pointer for this token (use actual stride, not assumed contiguous) + const CudaType* __restrict__ token_src = + reinterpret_cast(src) + token_idx * src_stride; + + // Destination bases in data and scale caches for this token's block + uint8_t* __restrict__ data_block = + data_cache + block_idx * data_block_stride; + uint8_t* __restrict__ scale_block = + sc_cache + block_idx * scale_block_stride; + + for (int g = tg_id; g < total_groups; g += num_thread_groups) { + const int head = g / groups_per_head; + const int group_in_head = g % groups_per_head; + + // Load 16 (or 8) bf16 elements from source + PVec in_vec; + const CudaType* __restrict__ src_ptr = + token_src + head * head_size + group_in_head * CVT_FP4_SF_VEC_SIZE + + tg_lane * ELTS; + +#pragma unroll + for (int i = 0; i < ELTS / 2; i++) { + in_vec.elts[i] = reinterpret_cast< + const typename PackedTypeConverter::Type*>(src_ptr)[i]; + } + + // Quantize: produces packed fp4 and writes scale factor. + uint8_t sf_val; + uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr; + + fp4_packed_t packed = cvt_warp_fp16_to_fp4( + in_vec, global_scale, sf_out_ptr); + + // Write packed FP4 data to data cache + uint8_t* __restrict__ data_dst = data_block + head * data_head_stride + + block_offset * data_block_offset_stride; + +#if CVT_FP4_PACK16 + { + // 16 elements → 8 bytes (u32x2) + int data_byte_offset = group_in_head * 8; + reinterpret_cast(data_dst + data_byte_offset)[0] = + (uint64_t(packed.hi) << 32) | uint64_t(packed.lo); + } +#else + { + // 8 elements → 4 bytes (uint32_t) + int data_byte_offset = + group_in_head * CVT_FP4_SF_VEC_SIZE / 2 + tg_lane * ELTS / 2; + reinterpret_cast(data_dst + data_byte_offset)[0] = packed; + } +#endif + + // Write block scale to scale cache. + // K (kv==0): linear layout (no swizzle). + // V (kv==1): swizzled layout for SM100 trtllm-gen MHA kernel. + if (sf_out_ptr != nullptr) { + int scale_idx = group_in_head; + uint8_t* __restrict__ scale_dst; + if (kv == 0) { + scale_dst = scale_block + head * scale_head_stride + + block_offset * scale_block_offset_stride + scale_idx; + } else { + int swizzled_offset = + swizzle_scale_offset(block_offset, scale_idx, scale_dim); + int swizzled_t = swizzled_offset / scale_dim; + int swizzled_s = swizzled_offset % scale_dim; + scale_dst = scale_block + head * scale_head_stride + + swizzled_t * scale_block_offset_stride + swizzled_s; + } + *scale_dst = sf_val; + } + } + } +} + +} // namespace vllm + +// Non-template entry point callable from cache_kernels.cu. +// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1]. +// Each KV side contains both data and scale: +// page = [K_data | K_scale | V_data | V_scale] +void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + torch::Tensor& slot_mapping, + torch::Tensor& k_scale, + torch::Tensor& v_scale) { + int num_tokens = slot_mapping.size(0); + int num_heads = key.size(1); + int head_size = key.size(2); + int data_dim = head_size / 2; + int scale_dim = head_size / 16; + int full_dim = data_dim + scale_dim; + + // key_cache is kv_cache[:, 0] with shape + // [num_blocks, block_size, num_heads, full_dim] in logical order. + // Strides encode the physical layout (HND or NHD). + TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D"); + TORCH_CHECK(key_cache.size(3) == full_dim, + "key_cache last dim must be data_dim + scale_dim, got ", + key_cache.size(3), " expected ", full_dim); + + int block_size = key_cache.size(1); + + TORCH_CHECK(head_size % 16 == 0, + "head_size must be divisible by 16 for NVFP4 KV cache"); + TORCH_CHECK(block_size % 4 == 0, + "block_size must be divisible by 4 for NVFP4 KV cache swizzle"); + + // Detect physical layout from strides (based on full_dim). + // HND: head stride > block_offset stride. + bool is_hnd = key_cache.stride(2) > key_cache.stride(1); + + int64_t data_block_stride = key_cache.stride(0); // page_bytes + int64_t data_head_stride, data_block_offset_stride; + if (is_hnd) { + data_head_stride = (int64_t)block_size * data_dim; + data_block_offset_stride = data_dim; + } else { + data_head_stride = data_dim; + data_block_offset_stride = (int64_t)num_heads * data_dim; + } + + // Page layout: [K_data | K_scale | V_data | V_scale] + // Scale follows data within each KV side. + int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim; + + uint8_t* key_scale_ptr = key_cache.data_ptr() + data_per_kv; + uint8_t* value_scale_ptr = value_cache.data_ptr() + data_per_kv; + + // Scale strides: same page stride, inner strides from layout. + int64_t scale_block_stride = data_block_stride; + int64_t scale_head_stride, scale_block_offset_stride; + if (is_hnd) { + scale_head_stride = (int64_t)block_size * scale_dim; + scale_block_offset_stride = scale_dim; + } else { + scale_head_stride = scale_dim; + scale_block_offset_stride = (int64_t)num_heads * scale_dim; + } + + const float* k_scale_ptr = k_scale.data_ptr(); + const float* v_scale_ptr = v_scale.data_ptr(); + + int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE; + int total_groups = num_heads * groups_per_head; + constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD; + int num_threads = std::min(total_groups * THREADS_PER_SF, 512); + num_threads = ((num_threads + 31) / 32) * 32; + + dim3 grid(num_tokens); + dim3 block(num_threads); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_REDUCED_FLOATING_TYPES( + key.scalar_type(), "reshape_and_cache_nvfp4", [&] { + vllm::reshape_and_cache_nvfp4_kernel + <<>>( + key.data_ptr(), value.data_ptr(), + key_cache.data_ptr(), value_cache.data_ptr(), + key_scale_ptr, value_scale_ptr, + slot_mapping.data_ptr(), k_scale_ptr, v_scale_ptr, + key.stride(0), value.stride(0), num_heads, head_size, + block_size, data_block_stride, data_head_stride, + data_block_offset_stride, scale_block_stride, scale_head_stride, + scale_block_offset_stride); + }); +} diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 0249461dd2f..9b022a042c8 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -10,7 +10,7 @@ from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize from vllm.platforms import current_platform -from vllm.utils.torch_utils import set_random_seed +from vllm.utils.torch_utils import nvfp4_kv_cache_split_views, set_random_seed COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")] DTYPES = [torch.bfloat16, torch.float] @@ -172,7 +172,7 @@ def test_reshape_and_cache( @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("seed", SEEDS) @pytest.mark.parametrize("device", CUDA_DEVICES) -@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE) +@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE + ["nvfp4"]) @pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS) @pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES) @pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS) @@ -202,6 +202,25 @@ def test_reshape_and_cache_flash( if kv_scale_type == "attn_head" and implementation != "cuda": pytest.skip("Only CUDA implementation supports attn_head scaling.") + if kv_cache_dtype == "nvfp4": + if not current_platform.has_device_capability(100): + pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).") + if implementation != "cuda": + pytest.skip("NVFP4 only supports CUDA implementation.") + if kv_scale_type != "tensor": + pytest.skip("NVFP4 only supports per-tensor scaling.") + if head_size % 16 != 0: + pytest.skip("NVFP4 requires head_size divisible by 16.") + if (head_size // 16) % 4 != 0: + pytest.skip( + "NVFP4 requires (head_size // 16) divisible by 4 " + "for 4x4 block scale swizzle." + ) + if block_size % 4 != 0: + pytest.skip("NVFP4 requires block_size divisible by 4.") + if dtype not in (torch.float16, torch.bfloat16): + pytest.skip("NVFP4 quantization only supports fp16/bf16 input.") + # fp8 conversion requires continugous memory buffer. Reduce the number of # blocks and tokens to consume less memory. num_tokens = num_tokens // 2 @@ -229,7 +248,23 @@ def test_reshape_and_cache_flash( del key_caches del value_caches - if kv_scale_type == "tensor": + # For nvfp4, the factory returns kv[:, 0] and kv[:, 1] like all dtypes. + # Split views are still needed for dequant verification. + key_scale_cache = None + value_scale_cache = None + nvfp4_key_data = None + nvfp4_value_data = None + if kv_cache_dtype == "nvfp4": + (nvfp4_key_data,), (key_scale_cache,) = nvfp4_kv_cache_split_views(key_cache) + (nvfp4_value_data,), (value_scale_cache,) = nvfp4_kv_cache_split_views( + value_cache + ) + + if kv_cache_dtype == "nvfp4": + # Global scale = amax / 448 (per-tensor) + k_scale = (key.abs().amax() / 448.0).to(torch.float32) + v_scale = (value.abs().amax() / 448.0).to(torch.float32) + elif kv_scale_type == "tensor": k_scale = (key.amax() / 64.0).to(torch.float32) v_scale = (value.amax() / 64.0).to(torch.float32) else: # "attn_head" @@ -240,8 +275,9 @@ def test_reshape_and_cache_flash( y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3) return y.contiguous() - key_cache_compact = permute_and_compact(key_cache) - value_cache_compact = permute_and_compact(value_cache) + if kv_cache_dtype != "nvfp4": + key_cache_compact = permute_and_compact(key_cache) + value_cache_compact = permute_and_compact(value_cache) def convert_fp8_local(output, input, scale, kv_dtype): fp8_input = input.view(current_platform.fp8_dtype()) @@ -257,7 +293,7 @@ def test_reshape_and_cache_flash( result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1) output.copy_(result) - # Clone the KV caches. + # Clone the KV caches (for non-nvfp4, used as reference baseline). if kv_cache_dtype == "fp8": cloned_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16) convert_fp8_local(cloned_key_cache, key_cache_compact, k_scale, kv_cache_dtype) @@ -265,25 +301,27 @@ def test_reshape_and_cache_flash( convert_fp8_local( cloned_value_cache, value_cache_compact, v_scale, kv_cache_dtype ) - else: + elif kv_cache_dtype != "nvfp4": cloned_key_cache = key_cache_compact.clone() cloned_value_cache = value_cache_compact.clone() + # Call the reshape_and_cache kernel. if implementation == "cuda": - opcheck( - torch.ops._C_cache_ops.reshape_and_cache_flash, - ( - key, - value, - key_cache, - value_cache, - slot_mapping, - kv_cache_dtype, - k_scale, - v_scale, - ), - cond=(head_size == HEAD_SIZES[0]), - ) + if kv_cache_dtype != "nvfp4": + opcheck( + torch.ops._C_cache_ops.reshape_and_cache_flash, + ( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ), + cond=(head_size == HEAD_SIZES[0]), + ) ops.reshape_and_cache_flash( key, value, @@ -309,6 +347,46 @@ def test_reshape_and_cache_flash( k_scale, v_scale, ) + + if kv_cache_dtype == "nvfp4": + # Verify NVFP4 by dequantizing the entire cache and comparing + # the written positions against original bf16 values. + # Same pattern as FP8: dequant whole cache, then extract and compare. + from tests.kernels.quantization.nvfp4_utils import ( + dequant_nvfp4_kv_cache, + ) + + def dequant_nvfp4_cache_nhd(data_cache, scale_cache, global_scale): + # data_cache: [N, T, H, data_dim] NHD (contiguous inner dims) + # scale_cache: [N, T, H, scale_dim] NHD (contiguous inner dims) + # Permute to HND layout for the dequant utility. + data_hnd = data_cache.permute(0, 2, 1, 3) + scale_hnd = scale_cache.permute(0, 2, 1, 3) + result_hnd = dequant_nvfp4_kv_cache( + data_hnd, scale_hnd, global_scale, head_size, block_size + ) + return result_hnd.permute(0, 2, 1, 3) # back to [N, T, H, D] + + result_key_cache = dequant_nvfp4_cache_nhd( + nvfp4_key_data, key_scale_cache, k_scale.item() + ) + result_value_cache = dequant_nvfp4_cache_nhd( + nvfp4_value_data, value_scale_cache, v_scale.item() + ) + + # Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping. + num_slots = num_blocks * block_size + result_key_flat = result_key_cache.reshape(num_slots, num_heads, head_size) + result_value_flat = result_value_cache.reshape(num_slots, num_heads, head_size) + + torch.testing.assert_close( + result_key_flat[slot_mapping], key.float(), atol=1.5, rtol=0.5 + ) + torch.testing.assert_close( + result_value_flat[slot_mapping], value.float(), atol=1.5, rtol=0.5 + ) + return + key_cache_compact = permute_and_compact(key_cache) value_cache_compact = permute_and_compact(value_cache) diff --git a/tests/kernels/quantization/nvfp4_utils.py b/tests/kernels/quantization/nvfp4_utils.py index 77889527143..df6513c131d 100644 --- a/tests/kernels/quantization/nvfp4_utils.py +++ b/tests/kernels/quantization/nvfp4_utils.py @@ -88,6 +88,60 @@ def break_fp4_bytes(a, dtype): return values.reshape(m, n * 2).to(dtype=dtype) +def dequant_nvfp4_kv_cache( + fp4_data: torch.Tensor, + block_scale: torch.Tensor, + global_scale: float, + head_size: int, + block_size: int, +) -> torch.Tensor: + """Dequantize an NVFP4 KV cache with 4x4-swizzled block scales. + + The input must be in HND layout so that the last two dims are + (block_size, last_dim). For NHD caches, permute to HND first. + + Args: + fp4_data: [..., num_heads, block_size, head_size//2] uint8 packed fp4. + block_scale: [..., num_heads, block_size, head_size//16] fp8 block + scales (as uint8 or float8_e4m3fn). + global_scale: checkpoint dequant scale (k_scale or v_scale). + head_size: head dimension. + block_size: page size. + + Returns: + [..., num_heads, block_size, head_size] float32. + """ + data_dim = head_size // 2 + scale_dim = head_size // 16 + + fp4_packed = fp4_data + sf_swizzled = block_scale.view(torch.uint8) + + # Unswizzle 4x4 block scales on (block_size, scale_dim) plane. + # [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S] + batch_shape = sf_swizzled.shape[:-2] + T, S = block_size, scale_dim + sg = S // 4 + sf_reshape = sf_swizzled.reshape(*batch_shape, T // 4, 4, sg, 4) + ndim = sf_reshape.ndim + # Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg) + perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2] + sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S) + sf_f32 = sf_linear.view(torch.float8_e4m3fn).to(torch.float32) + + # Unpack fp4 + shape = fp4_packed.shape # [..., T, data_dim] + fp4_flat = fp4_packed.reshape(-1, data_dim) + fp4_vals = break_fp4_bytes(fp4_flat, torch.float32) + fp4_vals = fp4_vals.reshape(*shape[:-1], head_size) + + # Dequant: fp4_val * block_scale * global_scale per 16-element group + return ( + fp4_vals.reshape(*shape[:-1], scale_dim, 16) + * (sf_f32 * global_scale).unsqueeze(-1) + ).reshape(*shape[:-1], head_size) + + def get_nvfp4_global_scale(a: torch.Tensor): return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(torch.float32) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index b84df7773c6..e34f57dea47 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -30,6 +30,7 @@ CacheDType = Literal[ "turboquant_3bit_nc", "int8_per_token_head", "fp8_per_token_head", + "nvfp4", ] MambaDType = Literal["auto", "float32", "float16"] MambaCacheMode = Literal["all", "align", "none"] diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index a92e2f4ad18..11985d2fa43 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -387,7 +387,9 @@ class Attention(nn.Module, AttentionLayerBase): self.query_quant = None if ( self.impl.supports_quant_query_input - and self.kv_cache_dtype.startswith("fp8") + and ( + self.kv_cache_dtype.startswith("fp8") or self.kv_cache_dtype == "nvfp4" + ) and not self.kv_cache_dtype.endswith("per_token_head") ): is_per_head = ( @@ -492,7 +494,7 @@ class Attention(nn.Module, AttentionLayerBase): # which reduces overheads during decoding. # Otherwise queries are quantized using custom ops # which causes decoding overheads - assert self.kv_cache_dtype in {"fp8", "fp8_e4m3"} + assert self.kv_cache_dtype in {"fp8", "fp8_e4m3", "nvfp4"} # check if query quantization is supported if self.impl.supports_quant_query_input: diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 26e377de69c..1eb9306ed4b 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -46,6 +46,7 @@ STR_DTYPE_TO_TORCH_DTYPE = { "turboquant_4bit_nc": torch.uint8, "turboquant_k3v4_nc": torch.uint8, "turboquant_3bit_nc": torch.uint8, + "nvfp4": torch.uint8, } TORCH_DTYPE_TO_NUMPY_DTYPE = { @@ -59,17 +60,19 @@ TORCH_DTYPE_TO_NUMPY_DTYPE = { MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = { - # TODO: Add more modelopt kv cache dtype - # mappings here when it supported by some attention backend - # (for example supports nvfp4). "fp8": "fp8_e4m3", + "nvfp4": "nvfp4", } T = TypeVar("T") def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: - return kv_cache_dtype.startswith("fp8") or kv_cache_dtype.endswith("per_token_head") + return ( + kv_cache_dtype.startswith("fp8") + or kv_cache_dtype.endswith("per_token_head") + or kv_cache_dtype == "nvfp4" + ) def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool: @@ -299,6 +302,8 @@ def get_kv_cache_quant_algo_string(quant_cfg: dict[str, Any]) -> str | None: and kv_algo.get("type") == "float" ): kv_algo = "fp8" + elif kv_algo.get("num_bits") == 4 and kv_algo.get("type") == "float": + kv_algo = "nvfp4" else: # Unknown/unsupported format - return "auto" as safe fallback logger.warning( @@ -375,6 +380,95 @@ def set_random_seed(seed: int | None) -> None: current_platform.manual_seed_all(seed) +def nvfp4_kv_cache_full_dim(head_size: int) -> int: + """Packed last dim for NVFP4 KV cache: fp4 data + fp8 block scales.""" + return head_size // 2 + head_size // 16 + + +def _nvfp4_split_data_scale( + kv_side: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Split a single NVFP4 KV-side buffer into data and scale views. + + The input is a 4D tensor for one KV side (K or V) whose last + dimension is ``full_dim = data_dim + scale_dim``. The physical + layout within each side is [data | scale], both packed contiguously. + + Args: + kv_side: 4D uint8 tensor with shape + ``(num_pages, dim_1, dim_2, full_dim)``. + May be in any permutation order (NHD or HND). + + Returns: + ``(data, scale)`` where + ``data`` is a uint8 view with shape + ``(num_pages, dim_1, dim_2, data_dim)``. + ``scale`` is a float8_e4m3fn view with shape + ``(num_pages, dim_1, dim_2, scale_dim)``. + """ + num_pages = kv_side.shape[0] + dim_1, dim_2 = kv_side.shape[1], kv_side.shape[2] + full_dim = kv_side.shape[3] + data_dim = full_dim * 8 // 9 + scale_dim = full_dim - data_dim + + data_per_kv = dim_1 * dim_2 * data_dim + page_bytes = kv_side.stride(0) + + # Derive inner strides from the kv_side strides, scaling by the + # ratio of the target dim to full_dim. This preserves the physical + # layout (NHD vs HND) encoded in the input tensor's strides. + s1 = kv_side.stride(1) * data_dim // full_dim + s2 = kv_side.stride(2) * data_dim // full_dim + data_shape = (num_pages, dim_1, dim_2, data_dim) + data_strides = (page_bytes, s1, s2, 1) + + s1_s = kv_side.stride(1) * scale_dim // full_dim + s2_s = kv_side.stride(2) * scale_dim // full_dim + scale_shape = (num_pages, dim_1, dim_2, scale_dim) + scale_strides = (page_bytes, s1_s, s2_s, 1) + + base = kv_side.storage_offset() + data = torch.as_strided(kv_side, data_shape, data_strides, storage_offset=base) + scale = torch.as_strided( + kv_side, scale_shape, scale_strides, storage_offset=base + data_per_kv + ).view(torch.float8_e4m3fn) + + return data, scale + + +def nvfp4_kv_cache_split_views(kv_cache: torch.Tensor) -> tuple[tuple, tuple]: + """Split an NVFP4 KV cache tensor into data and scale views. + + Accepts either a 5D tensor ``(num_pages, 2, dim_2, dim_3, full_dim)`` + or a 4D single-side tensor ``(num_pages, dim_2, dim_3, full_dim)``. + + Per-page layout: [K_data | K_scale | V_data | V_scale]. + Each KV side is self-contained (data followed by its scale), so the + 5D case simply splits each side independently. + + The returned views are in the same dim order as the input (NHD or + HND), so callers get views matching whichever order they passed in. + + Args: + kv_cache: 5D or 4D uint8 tensor where the last dimension is + ``full_dim = data_dim + scale_dim = 9 * head_size / 16``. + + Returns: + For 5D input: + ``(k_data, v_data), (k_scale, v_scale)`` + For 4D input (single KV side): + ``(data,), (scale,)`` + """ + if kv_cache.dim() == 4: + data, scale = _nvfp4_split_data_scale(kv_cache) + return (data,), (scale,) + + k_data, k_scale = _nvfp4_split_data_scale(kv_cache[:, 0]) + v_data, v_scale = _nvfp4_split_data_scale(kv_cache[:, 1]) + return (k_data, v_data), (k_scale, v_scale) + + def create_kv_caches_with_random_flash( num_blocks: int, block_size: int, @@ -401,15 +495,31 @@ def create_kv_caches_with_random_flash( value_caches: list[torch.Tensor] = [] for _ in range(num_layers): - key_value_cache = torch.empty( - size=kv_cache_allocation_shape, dtype=dtype, device=device - ).permute(*stride_order) - if cache_dtype in ["auto", "half", "bfloat16", "float"]: - key_value_cache.uniform_(-scale, scale) - elif cache_dtype == "fp8": - _generate_random_fp8(key_value_cache, -scale, scale) + if cache_dtype == "nvfp4": + # Full page dim: fp4 data + fp8 block scales per head. + # Per page layout: [K_data | K_scale | V_data | V_scale] + # Returns [:, 0] and [:, 1] like all other dtypes. + full_dim = nvfp4_kv_cache_full_dim(head_size) + nvfp4_shape = (num_blocks, 2, block_size, num_heads, full_dim) + nvfp4_phys = tuple(nvfp4_shape[i] for i in stride_order) + inv = [stride_order.index(i) for i in range(len(stride_order))] + key_value_cache = torch.randint( + 0, + 256, + nvfp4_phys, + dtype=dtype, + device=device, + ).permute(*inv) else: - raise ValueError(f"Does not support key cache of type {cache_dtype}") + key_value_cache = torch.empty( + size=kv_cache_allocation_shape, dtype=dtype, device=device + ).permute(*stride_order) + if cache_dtype in ["auto", "half", "bfloat16", "float"]: + key_value_cache.uniform_(-scale, scale) + elif cache_dtype == "fp8": + _generate_random_fp8(key_value_cache, -scale, scale) + else: + raise ValueError(f"Does not support key cache of type {cache_dtype}") key_caches.append(key_value_cache[:, 0]) value_caches.append(key_value_cache[:, 1]) return key_caches, value_caches diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 4adad61c2fa..78706e40c11 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -42,7 +42,12 @@ from vllm.utils.flashinfer import ( ) from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import is_pin_memory_available -from vllm.utils.torch_utils import is_quantized_kv_cache, is_strictly_contiguous +from vllm.utils.torch_utils import ( + is_quantized_kv_cache, + is_strictly_contiguous, + nvfp4_kv_cache_full_dim, + nvfp4_kv_cache_split_views, +) from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -355,6 +360,10 @@ class FlashInferBackend(AttentionBackend): head_size: int, cache_dtype_str: str = "auto", ) -> tuple[int, ...]: + if cache_dtype_str == "nvfp4": + # Packed layout: fp4 data + fp8 block scales in last dim + last_dim = nvfp4_kv_cache_full_dim(head_size) + return (num_blocks, 2, block_size, num_kv_heads, last_dim) return (num_blocks, 2, block_size, num_kv_heads, head_size) @staticmethod @@ -608,11 +617,19 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): self.cache_dtype = self.cache_config.cache_dtype # Cannot use self.kv_cache_spec.dtype here because kv_cache_spec # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3) - self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer( - self.cache_dtype - ) + self.is_kvcache_nvfp4 = self.cache_dtype == "nvfp4" + if self.is_kvcache_nvfp4: + # For NVFP4, kv_cache_dtype stays as the string "nvfp4" + # which is passed to FlashInferImpl + self.kv_cache_dtype = self.cache_dtype + raise NotImplementedError("nvfp4 KV cache is not yet supported") + else: + self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer( + self.cache_dtype + ) else: self.cache_dtype = "auto" + self.is_kvcache_nvfp4 = False assert self.kv_cache_spec.dtype == self.model_config.dtype self.kv_cache_dtype = self.kv_cache_spec.dtype @@ -626,7 +643,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): can_use_trtllm and not vllm_config.attention_config.disable_flashinfer_q_quantization ): - self.q_data_type = self.kv_cache_dtype + if self.is_kvcache_nvfp4: + # NVFP4 KV cache uses FP8 quantized queries + self.q_data_type = FlashInferBackend.get_fp8_dtype_for_flashinfer( + "fp8_e4m3" + ) + else: + self.q_data_type = self.kv_cache_dtype else: self.q_data_type = self.model_config.dtype @@ -1228,6 +1251,8 @@ class FlashInferImpl(AttentionImpl): self.sliding_window[0] if self.sliding_window is not None else -1 ) self.kv_cache_dtype = kv_cache_dtype + self.is_kvcache_nvfp4 = kv_cache_dtype == "nvfp4" + self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0 self.logits_soft_cap = logits_soft_cap self.kv_sharing_target_layer_name = kv_sharing_target_layer_name @@ -1406,7 +1431,16 @@ class FlashInferImpl(AttentionImpl): num_prefill_tokens = attn_metadata.num_prefill_tokens stride_order = FlashInferBackend.get_kv_cache_stride_order() - kv_cache_permute = kv_cache.permute(*stride_order) + kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous + + # For NVFP4, the kv_cache last dim is full_dim (data + scale packed). + # Split into correctly-strided data and scale views. + nvfp4_kv_data = None + nvfp4_kv_block_scales = None + if self.is_kvcache_nvfp4: + nvfp4_kv_data, nvfp4_kv_block_scales = nvfp4_kv_cache_split_views( + kv_cache_permute + ) use_dcp = self.dcp_world_size > 1 @@ -1490,8 +1524,20 @@ class FlashInferImpl(AttentionImpl): assert self.o_sf_scale is None out = output[num_decode_tokens:] - if attn_metadata.q_data_type != FP8_DTYPE and is_quantized_kv_cache( - self.kv_cache_dtype + prefill_kv_block_scales = None + if self.is_kvcache_nvfp4: + # NVFP4 trtllm-gen kernel requires FP8 query. + assert attn_metadata.q_data_type == FP8_DTYPE, ( + "NVFP4 KV cache requires FP8 quantized queries for " + "trtllm-gen prefill. Set " + "disable_flashinfer_q_quantization=False." + ) + mock_kv_cache = nvfp4_kv_data + mock_block_table = block_tables_prefill + prefill_kv_block_scales = nvfp4_kv_block_scales # noqa: F841 + elif ( + attn_metadata.q_data_type != FP8_DTYPE + and self.kv_cache_dtype.startswith("fp8") ): # TRTLLM prefill attention does not support BF16 Q # and fp8 kv cache. So to enable prefill attention @@ -1636,7 +1682,9 @@ class FlashInferImpl(AttentionImpl): trtllm_batch_decode_with_kv_cache( query=decode_query, - kv_cache=kv_cache_permute, + kv_cache=nvfp4_kv_data + if self.is_kvcache_nvfp4 + else kv_cache_permute, workspace_buffer=workspace_buffer, block_tables=block_tables_decode, seq_lens=seq_lens_decode, @@ -1667,11 +1715,13 @@ class FlashInferImpl(AttentionImpl): # and value[:num_actual_tokens] because the reshape_and_cache_flash # op uses the slot_mapping's shape to determine the number of # actual tokens. + k_cache = kv_cache[:, 0] + v_cache = kv_cache[:, 1] torch.ops._C_cache_ops.reshape_and_cache_flash( key, value, - kv_cache[:, 0], - kv_cache[:, 1], + k_cache, + v_cache, slot_mapping, self.kv_cache_dtype, layer._k_scale, diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 8aed95ddd0e..05a6eba824a 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -17,7 +17,7 @@ from vllm.logger import init_logger if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import get_dtype_size +from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim logger = init_logger(__name__) @@ -38,11 +38,20 @@ class KVQuantMode(IntEnum): FP8_PER_TENSOR = 1 # per-tensor scales (current fp8 path) INT8_PER_TOKEN_HEAD = 2 # per-token-head dynamic scales for int8 FP8_PER_TOKEN_HEAD = 3 # per-token-head dynamic scales for fp8 + NVFP4 = 4 # packed fp4 data + fp8 block scales @property def is_per_token_head(self) -> bool: """True for any per-token-head quantization mode.""" - return self >= 2 + return self in ( + KVQuantMode.INT8_PER_TOKEN_HEAD, + KVQuantMode.FP8_PER_TOKEN_HEAD, + ) + + @property + def is_nvfp4(self) -> bool: + """True for NVFP4 packed quantization mode.""" + return self == KVQuantMode.NVFP4 def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: @@ -51,7 +60,9 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: return KVQuantMode.INT8_PER_TOKEN_HEAD if kv_cache_dtype == "fp8_per_token_head": return KVQuantMode.FP8_PER_TOKEN_HEAD - if kv_cache_dtype.startswith("fp8"): + if kv_cache_dtype == "nvfp4": + return KVQuantMode.NVFP4 + if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"): return KVQuantMode.FP8_PER_TENSOR return KVQuantMode.NONE @@ -237,6 +248,19 @@ class FullAttentionSpec(AttentionSpec): @property def real_page_size_bytes(self) -> int: + if self.kv_quant_mode.is_nvfp4: + # Packed layout per head: fp4 data + fp8 block scales. + # fp4 data: head_size//2 bytes (2 fp4 values per byte) + # fp8 block scale: head_size//16 bytes (1 scale per 16 elements) + last_dim = nvfp4_kv_cache_full_dim( + self.head_size + ) + nvfp4_kv_cache_full_dim(self.head_size_v) + return ( + self.block_size + * self.num_kv_heads + * last_dim + * get_dtype_size(self.dtype) + ) return ( self.block_size * self.num_kv_heads From 1174723eba176fcc68673a5a0a2c0090416aada8 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Fri, 17 Apr 2026 10:31:41 -0400 Subject: [PATCH 087/150] Fix TURBOQUANT backend selection in cuda.py (#40060) Signed-off-by: Michael Goin --- docs/design/attention_backends.md | 2 ++ vllm/platforms/cuda.py | 7 ++----- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 65e93657e78..d10c23038a6 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -106,6 +106,7 @@ Priority is **1 = highest** (tried first). | 2 | `FLASH_ATTN` | | 3 | `TRITON_ATTN` | | 4 | `FLEX_ATTENTION` | +| 5 | `TURBOQUANT` | **Ampere/Hopper (SM 8.x-9.x):** @@ -115,6 +116,7 @@ Priority is **1 = highest** (tried first). | 2 | `FLASHINFER` | | 3 | `TRITON_ATTN` | | 4 | `FLEX_ATTENTION` | +| 5 | `TURBOQUANT` | ### MLA Attention (DeepSeek-style) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index bbbd9af0cf7..d79d3191820 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -131,6 +131,7 @@ def _get_backend_priorities( AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TRITON_ATTN, AttentionBackendEnum.FLEX_ATTENTION, + AttentionBackendEnum.TURBOQUANT, ] else: return [ @@ -138,6 +139,7 @@ def _get_backend_priorities( AttentionBackendEnum.FLASHINFER, AttentionBackendEnum.TRITON_ATTN, AttentionBackendEnum.FLEX_ATTENTION, + AttentionBackendEnum.TURBOQUANT, ] @@ -255,11 +257,6 @@ class CudaPlatformBase(Platform): valid_backends_priorities = [] invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {} - # TurboQuant KV cache: route directly to TQ backend - kv_cache_dtype = attn_selector_config.kv_cache_dtype - if kv_cache_dtype is not None and kv_cache_dtype.startswith("turboquant_"): - return [(AttentionBackendEnum.TURBOQUANT, 0)], {} - backend_priorities = _get_backend_priorities( attn_selector_config.use_mla, device_capability, From 747256bb5d645f28579bdb1a913f3584b6a246f6 Mon Sep 17 00:00:00 2001 From: Jing Wang Date: Sat, 18 Apr 2026 00:02:50 +0800 Subject: [PATCH 088/150] [Bugfix][Core] Fix stuck chunked pipeline parallelism with async scheduling (#38726) Signed-off-by: Jing Wang Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- vllm/v1/worker/gpu_model_runner.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index d95e70d071a..1e05d68078e 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -3758,6 +3758,15 @@ class GPUModelRunner( return slot_mappings_by_gid, slot_mappings_by_layer + def _is_all_reqs_chunked_prefill(self) -> bool: + """Check if all scheduled requests are marked to discard sampled tokens. + + This is true when `discard_request_mask` is set for every scheduled + request (e.g., for chunked prefill requests that are not the last + prefill chunk).""" + num_reqs = self.input_batch.num_reqs + return bool(self.discard_request_mask.np[:num_reqs].all()) + @torch.inference_mode() def execute_model( self, @@ -4361,9 +4370,12 @@ class GPUModelRunner( assert sampled_token_ids.dim() == 2 and sampled_token_ids.shape[-1] == 1, ( "PP+async expects sampled_token_ids to have shape [num_reqs, 1]" ) - torch.distributed.broadcast( - sampled_token_ids, src=pp.rank, group=pp.device_group - ) + # Skip for chunked prefill: sampled tokens are dummy + # and will be discarded, no need to broadcast. + if not self._is_all_reqs_chunked_prefill(): + torch.distributed.broadcast( + sampled_token_ids, src=pp.rank, group=pp.device_group + ) def _pp_receive_prev_sampled_token_ids_to_input_batch(self) -> None: """Receive sampled token ids broadcast from last PP stage""" @@ -4372,7 +4384,9 @@ class GPUModelRunner( num_reqs = self.input_batch.num_reqs # `prev_sampled_token_ids` is expected to have shape [num_reqs, 1]. recv = torch.empty((num_reqs, 1), dtype=torch.int32, device=self.device) - torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group) + # skip for chunked prefill. + if not self._is_all_reqs_chunked_prefill(): + torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group) self.input_batch.prev_sampled_token_ids = recv # construct `prev_req_id_to_index` here so `_prepare_input_ids` From ceade1952c7aebfbf328daa7dd40877c5cd8e8a5 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Sat, 18 Apr 2026 00:38:10 +0800 Subject: [PATCH 089/150] [BugFix] Support custom tool parsers when tool_choice is `required` and named function (#39870) Signed-off-by: JaredforReal Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: sfeng33 <4florafeng@gmail.com> --- .../openai/chat_completion/serving.py | 45 ++++++++++++++++--- vllm/entrypoints/openai/engine/serving.py | 31 ++++++++++--- vllm/tool_parsers/abstract_tool_parser.py | 11 +++++ vllm/tool_parsers/glm47_moe_tool_parser.py | 2 + vllm/tool_parsers/glm4_moe_tool_parser.py | 23 +++++++++- 5 files changed, 100 insertions(+), 12 deletions(-) diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 446f127a91e..4a56c63ecd4 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -557,6 +557,20 @@ class OpenAIServingChat(OpenAIServing): and self._should_stream_with_auto_tool_parsing(request) ) + # Determine whether required/named tool_choice should fall back to + # the auto tool_parser path instead of the standard JSON-based parsing. + # This happens when the parser declares supports_required_and_named=False + # (e.g. GLM models that output XML instead of JSON). + tool_choice_uses_parser = ( + self.tool_parser is not None + and not self.tool_parser.supports_required_and_named + and request.tools + and ( + request.tool_choice == "required" + or isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) + ) + ) + all_previous_token_ids: list[list[int]] | None function_name_returned = [False] * num_choices if self.tool_call_id_type == "kimi_k2": @@ -569,7 +583,12 @@ class OpenAIServingChat(OpenAIServing): # Only one of these will be used, thus previous_texts and # all_previous_token_ids will not be used twice in the same iteration. - if is_mistral_grammar_path or tool_choice_auto or reasoning_parser: + if ( + is_mistral_grammar_path + or tool_choice_auto + or tool_choice_uses_parser + or reasoning_parser + ): # These are only required in "auto" tool choice case all_previous_token_ids = [[] for _ in range(num_choices)] reasoning_end_arr = [False] * num_choices @@ -764,7 +783,12 @@ class OpenAIServingChat(OpenAIServing): delta_message: DeltaMessage | None # just update previous_texts and previous_token_ids - if is_mistral_grammar_path or tool_choice_auto or reasoning_parser: + if ( + is_mistral_grammar_path + or tool_choice_auto + or tool_choice_uses_parser + or reasoning_parser + ): assert previous_texts is not None assert all_previous_token_ids is not None previous_text = previous_texts[i] @@ -813,7 +837,9 @@ class OpenAIServingChat(OpenAIServing): if result.tools_called: tools_streamed[i] = True # handle streaming deltas for tools with named tool_choice - elif tool_choice_function_name: + # Skip when tool_choice_uses_parser so it falls through + # to the auto tool_parser branches below. + elif tool_choice_function_name and not tool_choice_uses_parser: # When encountering think end id in prompt_token_ids # i.e {"enable_thinking": False}, # check BEFORE calling the parser to avoid a spurious @@ -851,7 +877,6 @@ class OpenAIServingChat(OpenAIServing): ): reasoning_end_arr[i] = True if delta_message and delta_message.content: - # This need to be added to next `delta_text` current_text = delta_message.content delta_message.content = None else: @@ -896,7 +921,12 @@ class OpenAIServingChat(OpenAIServing): ) tools_streamed[i] = True - elif request.tool_choice == "required": + # Skip when tool_choice_uses_parser so it falls through + # to the auto tool_parser branches below. + elif ( + request.tool_choice == "required" + and not tool_choice_uses_parser + ): assert previous_texts is not None previous_text = previous_texts[i] current_text = previous_text + delta_text @@ -966,7 +996,10 @@ class OpenAIServingChat(OpenAIServing): # update the previous values for the next iteration if ( - is_mistral_grammar_path or tool_choice_auto or reasoning_parser + is_mistral_grammar_path + or tool_choice_auto + or tool_choice_uses_parser + or reasoning_parser ) and not self.use_harmony: assert previous_texts is not None assert all_previous_token_ids is not None diff --git a/vllm/entrypoints/openai/engine/serving.py b/vllm/entrypoints/openai/engine/serving.py index 85237604e5a..ab33008aeeb 100644 --- a/vllm/entrypoints/openai/engine/serving.py +++ b/vllm/entrypoints/openai/engine/serving.py @@ -627,7 +627,7 @@ class OpenAIServing: and isinstance(request.tool_choice, ToolChoiceFunction) ): assert content is not None - # Forced Function Call + # Forced Function Call (Responses API) function_calls.append( FunctionCall(name=request.tool_choice.name, arguments=content) ) @@ -636,14 +636,20 @@ class OpenAIServing: not use_mistral_tool_parser and request.tool_choice and isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) + and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named) ): + # Named function with standard JSON-based parsing assert content is not None - # Forced Function Call function_calls.append( FunctionCall(name=request.tool_choice.function.name, arguments=content) ) content = None # Clear content since tool is called. - elif not use_mistral_tool_parser and request.tool_choice == "required": + elif ( + not use_mistral_tool_parser + and request.tool_choice == "required" + and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named) + ): + # "required" with standard JSON-based parsing tool_calls = [] with contextlib.suppress(ValidationError): content = content or "" @@ -662,15 +668,30 @@ class OpenAIServing: use_mistral_tool_parser or ( enable_auto_tools - and (request.tool_choice == "auto" or request.tool_choice is None) + and ( + request.tool_choice == "auto" + or request.tool_choice is None + or ( + not tool_parser_cls.supports_required_and_named + and request.tools + and ( + request.tool_choice == "required" + or isinstance( + request.tool_choice, + ChatCompletionNamedToolChoiceParam, + ) + ) + ) + ) ) ): + # Automatic Tool Call Parsing (also used as fallback for + # required/named when supports_required_and_named=False) if tokenizer is None: raise ValueError( "Tokenizer not available when `skip_tokenizer_init=True`" ) - # Automatic Tool Call Parsing try: tool_parser = tool_parser_cls(tokenizer, request.tools) except RuntimeError as e: diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index c127bd9dd68..bc0c090d8bf 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/vllm/tool_parsers/abstract_tool_parser.py @@ -44,6 +44,17 @@ class ToolParser: derived classes. """ + # When True (default), the serving layer uses the standard JSON-based + # parsing for tool_choice="required" and named function tool_choice, + # which works for models where guided decoding produces well-formed + # JSON output (e.g. Hermes). + # Subclasses set False when the standard parsing does not work for + # their model's output format (e.g. GLM models that use XML). When + # False, the serving layer falls back to the tool_parser's + # extract_tool_calls / extract_tool_calls_streaming methods for + # required/named tool_choice, treating them the same as "auto". + supports_required_and_named: bool = True + def __init__( self, tokenizer: TokenizerLike, diff --git a/vllm/tool_parsers/glm47_moe_tool_parser.py b/vllm/tool_parsers/glm47_moe_tool_parser.py index 765d6d37de1..47b6ad2f5af 100644 --- a/vllm/tool_parsers/glm47_moe_tool_parser.py +++ b/vllm/tool_parsers/glm47_moe_tool_parser.py @@ -23,6 +23,8 @@ logger = init_logger(__name__) class Glm47MoeModelToolParser(Glm4MoeModelToolParser): + supports_required_and_named = False + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) # GLM-4.7 format: func_name[...]* diff --git a/vllm/tool_parsers/glm4_moe_tool_parser.py b/vllm/tool_parsers/glm4_moe_tool_parser.py index 491c9599b61..ac266ede3f4 100644 --- a/vllm/tool_parsers/glm4_moe_tool_parser.py +++ b/vllm/tool_parsers/glm4_moe_tool_parser.py @@ -20,6 +20,7 @@ import regex as re from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ) from vllm.entrypoints.openai.engine.protocol import ( @@ -50,6 +51,8 @@ class Glm4MoeModelToolParser(ToolParser): call, and diffs against what was previously sent to emit only new content. """ + supports_required_and_named = False + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) # Stateful streaming fields @@ -156,7 +159,25 @@ class Glm4MoeModelToolParser(ToolParser): def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest ) -> ChatCompletionRequest | ResponsesRequest: - """Adjust request parameters for tool call token handling.""" + """Adjust request parameters for tool call token handling. + + For required/named tool_choice, skip setting structured_outputs + because GLM models output tool calls in XML format (per chat + template). Guided decoding would force JSON output, conflicting + with the XML format and causing parsing failures. + """ + if request.tools: + tc = request.tool_choice + if tc == "required" or isinstance(tc, ChatCompletionNamedToolChoiceParam): + # Do NOT call super().adjust_request() for required/named, + # because it would set structured_outputs and force JSON + # output via guided decoding. GLM models use XML tool-call + # syntax (defined in the chat template), so guided decoding + # must be skipped to let the model output XML freely. + # The tool_parser handles extraction from XML output. + if request.tool_choice != "none": + request.skip_special_tokens = False + return request request = super().adjust_request(request) if request.tools and request.tool_choice != "none": # Ensure tool call tokens (, ) are not skipped From 640cc9dd7dae3ba08f4dc6e479403fbaf99f2d93 Mon Sep 17 00:00:00 2001 From: allgather Date: Fri, 17 Apr 2026 09:39:19 -0700 Subject: [PATCH 090/150] feat: Add LoRA support for Gemma4ForConditionalGeneration (#39291) Signed-off-by: allgather Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/model_executor/models/gemma4_mm.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 08978ec3f44..ff1760dde08 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -67,6 +67,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( MultiModalEmbeddings, SupportsEagle3, + SupportsLoRA, SupportsMultiModal, SupportsPP, ) @@ -880,6 +881,7 @@ class Gemma4ForConditionalGeneration( nn.Module, SupportsMultiModal, SupportsPP, + SupportsLoRA, SupportsEagle3, ): packed_modules_mapping = { @@ -1358,10 +1360,16 @@ class Gemma4ForConditionalGeneration( def get_mm_mapping(self) -> MultiModelKeys: """Get the module prefix mapping for multimodal models.""" + connectors = ["embed_vision"] + tower_models = ["vision_tower"] + if self.audio_tower is not None: + connectors.append("embed_audio") + tower_models.append("audio_tower") + return MultiModelKeys.from_string_field( language_model="language_model", - connector=["embed_vision", "embed_audio"], - tower_model=["vision_tower", "audio_tower"], + connector=connectors, + tower_model=tower_models, ) @classmethod From 512765d52d743ebaefb7705ff775226d6a7fcc7a Mon Sep 17 00:00:00 2001 From: Roger Wang Date: Fri, 17 Apr 2026 09:49:21 -0700 Subject: [PATCH 091/150] [Misc][UX] Map mimo reasoning and tooling parsers (#40089) Signed-off-by: Roger Wang Co-authored-by: Chauncey --- vllm/reasoning/__init__.py | 4 ++++ vllm/tool_parsers/__init__.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index 2d57b93369d..37d8a9b1dab 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -60,6 +60,10 @@ _REASONING_PARSERS_TO_REGISTER = { "kimi_k2_reasoning_parser", "KimiK2ReasoningParser", ), + "mimo": ( + "qwen3_reasoning_parser", + "Qwen3ReasoningParser", + ), "minimax_m2": ( "minimax_m2_reasoning_parser", "MiniMaxM2ReasoningParser", diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index bffa00c4ef3..abfa8f3fdbe 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -94,6 +94,10 @@ _TOOL_PARSERS_TO_REGISTER = { "longcat_tool_parser", "LongcatFlashToolParser", ), + "mimo": ( + "qwen3xml_tool_parser", + "Qwen3XMLToolParser", + ), "minimax_m2": ( "minimax_m2_tool_parser", "MinimaxM2ToolParser", From 251c18d1f89c363b9f5ecaa29247df64f4157308 Mon Sep 17 00:00:00 2001 From: Xinyu Chen Date: Sat, 18 Apr 2026 00:55:08 +0800 Subject: [PATCH 092/150] skip fp8e4b15 on xpu (#39957) Signed-off-by: Xinyu Chen Co-authored-by: Kunshang Ji --- tests/quantization/test_turboquant.py | 32 ++++++++++--------- .../attention/ops/triton_turboquant_decode.py | 12 +++++-- 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/tests/quantization/test_turboquant.py b/tests/quantization/test_turboquant.py index 78c137e6762..519022346c2 100644 --- a/tests/quantization/test_turboquant.py +++ b/tests/quantization/test_turboquant.py @@ -21,6 +21,7 @@ from vllm.model_executor.layers.quantization.turboquant.config import ( from vllm.model_executor.layers.quantization.turboquant.quantizer import ( generate_wht_signs, ) +from vllm.platforms import current_platform from vllm.utils.math_utils import next_power_of_2 # ============================================================================ @@ -345,7 +346,8 @@ class TestLloydMax: # Rotation matrix tests (GPU required) # ============================================================================ -CUDA_AVAILABLE = torch.cuda.is_available() +GPGPU_AVAILABLE = torch.cuda.is_available() or torch.xpu.is_available() +DEVICE_TYPE = current_platform.device_type def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor: @@ -360,16 +362,16 @@ def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Te return Q.to(device) -@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available") +@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available") class TestRotationMatrix: """Tests for the QR-based rotation (standalone benchmarks only).""" @pytest.mark.parametrize("dim", [64, 96, 128, 256]) def test_rotation_matrix_shape_and_orthogonal(self, dim): - Pi = generate_rotation_matrix(dim, seed=42, device="cuda") + Pi = generate_rotation_matrix(dim, seed=42, device=DEVICE_TYPE) assert Pi.shape == (dim, dim) eye = Pi @ Pi.T - assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), ( + assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), ( f"Pi not orthogonal for dim={dim}" ) @@ -385,7 +387,7 @@ class TestRotationMatrix: def test_rotation_matrix_det_is_pm1(self): """Orthogonal matrix determinant must be +1 or -1.""" - Pi = generate_rotation_matrix(128, seed=42, device="cuda") + Pi = generate_rotation_matrix(128, seed=42, device=DEVICE_TYPE) det = torch.linalg.det(Pi) assert abs(abs(det.item()) - 1.0) < 1e-4 @@ -403,31 +405,31 @@ def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor: return (H / math.sqrt(d)).to(torch.device(device)) -@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available") +@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available") class TestWHTRotation: """Tests for the WHT rotation actually used in serving.""" @pytest.mark.parametrize("dim", [64, 128, 256]) def test_wht_orthonormal(self, dim): """signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I.""" - signs = generate_wht_signs(dim, seed=42, device="cuda") - H = _build_hadamard(dim, "cuda") + signs = generate_wht_signs(dim, seed=42, device=DEVICE_TYPE) + H = _build_hadamard(dim, DEVICE_TYPE) PiT = (signs.unsqueeze(1) * H).contiguous() eye = PiT @ PiT.T - assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), ( + assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), ( f"WHT rotation not orthonormal for dim={dim}" ) @pytest.mark.parametrize("dim", [64, 128, 256]) def test_wht_self_inverse(self, dim): """PiT should be self-inverse: PiT @ PiT = I (up to sign flip).""" - signs = generate_wht_signs(dim, seed=42, device="cuda") - H = _build_hadamard(dim, "cuda") + signs = generate_wht_signs(dim, seed=42, device=DEVICE_TYPE) + H = _build_hadamard(dim, DEVICE_TYPE) PiT = (signs.unsqueeze(1) * H).contiguous() Pi = PiT.T.contiguous() # Pi @ PiT should be identity (rotation then inverse) result = Pi @ PiT - assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), ( + assert torch.allclose(result, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), ( f"WHT rotation not self-inverse for dim={dim}" ) @@ -454,7 +456,7 @@ class TestWHTRotation: # ============================================================================ -@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available") +@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available") class TestStoreDecodeRoundTrip: """End-to-end: store KV into TQ cache, decode, compare vs fp16 ref.""" @@ -487,11 +489,11 @@ class TestStoreDecodeRoundTrip: block_size = 16 num_blocks = 1 - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) # Generate rotation signs = generate_wht_signs(D, seed=42, device=device) - H = _build_hadamard(D, "cuda") + H = _build_hadamard(D, DEVICE_TYPE) PiT = (signs.unsqueeze(1) * H).contiguous().float() Pi = PiT.T.contiguous() diff --git a/vllm/v1/attention/ops/triton_turboquant_decode.py b/vllm/v1/attention/ops/triton_turboquant_decode.py index 8b276e31eaf..8a6117afa8e 100644 --- a/vllm/v1/attention/ops/triton_turboquant_decode.py +++ b/vllm/v1/attention/ops/triton_turboquant_decode.py @@ -13,6 +13,7 @@ from typing import Any import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.v1.attention.ops.triton_decode_attention import ( _fwd_kernel_stage2, @@ -22,10 +23,15 @@ _FP8_E4B15: dict[int, int] = {} def _use_fp8_e4b15(device: int = 0) -> int: - """Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0.""" + """Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0. + On non-CUDA platforms (e.g. XPU), always returns 0 (use e4nv format). + """ if device not in _FP8_E4B15: - cap = torch.cuda.get_device_capability(device) - _FP8_E4B15[device] = 1 if cap < (8, 9) else 0 + if current_platform.is_cuda_alike(): + cap = torch.cuda.get_device_capability(device) + _FP8_E4B15[device] = 1 if cap < (8, 9) else 0 + else: + _FP8_E4B15[device] = 0 return _FP8_E4B15[device] From 1ae11e2bfcf948876487bf7319f5bbb049768ba4 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 17 Apr 2026 14:30:08 -0500 Subject: [PATCH 093/150] [ROCm][CI] Build fastsafetensors from source so it links against libamdhip64 (#39978) Signed-off-by: Andreas Karatzas --- docker/Dockerfile.rocm | 24 ------------------------ requirements/rocm.txt | 2 +- requirements/test/rocm.in | 2 +- requirements/test/rocm.txt | 2 +- 4 files changed, 3 insertions(+), 27 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index fa7a5846edc..fb4d04430b3 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -262,30 +262,6 @@ RUN --mount=type=bind,source=.git,target=vllm/.git \ && echo "Detected vLLM version: ${VLLM_VERSION}" \ && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt -# Fail if git-based package dependencies are found in requirements files -# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI) -# Extra notes: pip install is able to handle git+ URLs, but uv doesn't. -RUN echo "Checking for git-based packages in requirements files..." \ - && echo "Checking common.txt for git-based packages:" \ - && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \ - echo "ERROR: Git-based packages found in common.txt:"; \ - grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \ - echo "Please publish these packages to PyPI instead of using git dependencies."; \ - exit 1; \ - else \ - echo " ✓ No git-based packages found in common.txt"; \ - fi \ - && echo "Checking rocm.txt for git-based packages:" \ - && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \ - echo "ERROR: Git-based packages found in rocm.txt:"; \ - grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \ - echo "Please publish these packages to PyPI instead of using git dependencies."; \ - exit 1; \ - else \ - echo " ✓ No git-based packages found in rocm.txt"; \ - fi \ - && echo "All requirements files are clean - no git-based packages found" - # Pin vLLM dependencies to exact versions of custom ROCm wheels # This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py diff --git a/requirements/rocm.txt b/requirements/rocm.txt index 011d0e53fd1..deaeae2d5e2 100644 --- a/requirements/rocm.txt +++ b/requirements/rocm.txt @@ -22,4 +22,4 @@ timm>=1.0.17 # To be consistent with test_quark.py amd-quark>=0.8.99 # Required for faster safetensors model loading -fastsafetensors >= 0.2.2 \ No newline at end of file +fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 \ No newline at end of file diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 0c9335240f5..0ae2e487853 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -55,7 +55,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test numba==0.61.2 # Required for N-gram speculative decoding numpy runai-model-streamer[s3,gcs,azure]==0.15.7 -fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage +fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels instanttensor>=0.1.5 pydantic>=2.12 # 2.11 leads to error on python 3.13 decord==0.6.0 diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index eaa11451a21..5abec4dce4c 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -277,7 +277,7 @@ fastar==0.10.0 # via fastapi-cloud-cli fastparquet==2026.3.0 # via genai-perf -fastsafetensors==0.2.2 +fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c # via # -c requirements/rocm.txt # -r requirements/test/rocm.in From 58da4ee047a97b71776488301f66612838b79788 Mon Sep 17 00:00:00 2001 From: Ryan Rock Date: Fri, 17 Apr 2026 14:30:20 -0500 Subject: [PATCH 094/150] [AMD][CI] Update DeepEP branch (#38396) Signed-off-by: Ryan Rock --- .buildkite/test-amd.yaml | 2 +- docker/Dockerfile.rocm | 15 +++++++-------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 847d8762da2..fe617cb10a1 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2613,6 +2613,7 @@ steps: - vllm/platforms/rocm.py commands: - export TORCH_NCCL_BLOCKING_WAIT=1 + - VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - pytest -v -s tests/v1/distributed/test_dbo.py @@ -3601,7 +3602,6 @@ steps: commands: - export TORCH_NCCL_BLOCKING_WAIT=1 - pytest -v -s tests/distributed/test_context_parallel.py - - VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - pytest -v -s tests/v1/distributed/test_dbo.py diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index fb4d04430b3..d1eebcd7bd4 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -192,9 +192,10 @@ RUN cd /opt/rixl && mkdir -p /app/install && \ FROM base AS build_deep ARG ROCSHMEM_BRANCH="ba0bf0f3" ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" -ARG DEEPEP_BRANCH="e84464ec" +ARG DEEPEP_BRANCH="5d90af8b" ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" ARG DEEPEP_NIC="cx7" +ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" ENV ROCSHMEM_DIR=/opt/rocshmem RUN git clone ${ROCSHMEM_REPO} \ @@ -202,13 +203,11 @@ RUN git clone ${ROCSHMEM_REPO} \ && git checkout ${ROCSHMEM_BRANCH} \ && mkdir -p projects/rocshmem/build \ && cd projects/rocshmem/build \ - && cmake .. \ - -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \ - -DROCM_PATH=/opt/rocm \ - -DCMAKE_POSITION_INDEPENDENT_CODE=ON \ - -DUSE_EXTERNAL_MPI=OFF \ - && make -j \ - && make install + && bash ../scripts/build_configs/all_backends \ + -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \ + -DROCM_PATH=/opt/rocm \ + -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ + -DUSE_EXTERNAL_MPI=OFF # Build DeepEP wheel. # DeepEP looks for rocshmem at ROCSHMEM_DIR. From 6ef1efd51f11106fc44deb9e7b2f5cd1247fc37e Mon Sep 17 00:00:00 2001 From: aditi-amd Date: Fri, 17 Apr 2026 13:08:37 -0700 Subject: [PATCH 095/150] [ROCm] Fix TurboQuant on ROCm: backend routing, flash-attn compat, int64 overflow (#39953) Signed-off-by: aditi --- vllm/platforms/rocm.py | 1 + vllm/v1/attention/backends/turboquant_attn.py | 14 +++----------- .../attention/ops/triton_turboquant_decode.py | 12 ++++++------ .../attention/ops/triton_turboquant_store.py | 18 ++++++++++++------ 4 files changed, 22 insertions(+), 23 deletions(-) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 11a85e924b8..89714f00f64 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -382,6 +382,7 @@ def _get_backend_priorities( if is_aiter_found_and_supported(): backends.append(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN) backends.append(AttentionBackendEnum.TRITON_ATTN) + backends.append(AttentionBackendEnum.TURBOQUANT) return backends diff --git a/vllm/v1/attention/backends/turboquant_attn.py b/vllm/v1/attention/backends/turboquant_attn.py index 279fcb04ace..2659cbffa82 100644 --- a/vllm/v1/attention/backends/turboquant_attn.py +++ b/vllm/v1/attention/backends/turboquant_attn.py @@ -507,8 +507,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]): # max_query_len == max_seq_len means no request has prior cached KV. # Both are Python ints — no GPU sync. if _HAS_FLASH_ATTN and attn_metadata.max_query_len == attn_metadata.max_seq_len: - output = torch.empty(N, Hq, D, device=query.device, dtype=query.dtype) - flash_attn_varlen_func( + return flash_attn_varlen_func( q=query, k=key, v=value, @@ -518,9 +517,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]): max_seqlen_k=attn_metadata.max_query_len, softmax_scale=self.scale, causal=True, - out=output, ) - return output # Continuation or no flash_attn: per-request attention. # For continuation chunks (seq_len > q_len), we must attend to @@ -557,10 +554,9 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]): if q_len == seq_len: # First-chunk prefill: all K/V are in the current batch. if _HAS_FLASH_ATTN: - out = torch.empty_like(q_seq) _cu_2[1] = q_len cu = _cu_2 - flash_attn_varlen_func( + out = flash_attn_varlen_func( q=q_seq, k=k_seq, v=v_seq, @@ -570,7 +566,6 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]): max_seqlen_k=q_len, softmax_scale=self.scale, causal=True, - out=out, ) else: q_t = q_seq.transpose(0, 1).contiguous() @@ -733,10 +728,9 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]): # Attention: q_len queries attending to seq_len K/V with causal mask if _HAS_FLASH_ATTN: - output = torch.empty(q_len, Hq, D, device=device, dtype=query.dtype) cu_seqlens_q = torch.tensor([0, q_len], device=device, dtype=torch.int32) cu_seqlens_k = torch.tensor([0, seq_len], device=device, dtype=torch.int32) - flash_attn_varlen_func( + return flash_attn_varlen_func( q=query, k=k_full, v=v_full, @@ -746,9 +740,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]): max_seqlen_k=seq_len, softmax_scale=self.scale, causal=True, - out=output, ) - return output else: # SDPA fallback: expand KV for GQA, build causal mask q_t = query.transpose(0, 1).unsqueeze(0) # (1, Hq, q_len, D) diff --git a/vllm/v1/attention/ops/triton_turboquant_decode.py b/vllm/v1/attention/ops/triton_turboquant_decode.py index 8a6117afa8e..a789f9be7bb 100644 --- a/vllm/v1/attention/ops/triton_turboquant_decode.py +++ b/vllm/v1/attention/ops/triton_turboquant_decode.py @@ -143,12 +143,12 @@ def _tq_decode_stage1( Block_table_ptr + bt_base + page_idx, mask=kv_mask, other=0, - ) + ).to(tl.int64) slot_bases = ( block_nums * stride_cache_block - + page_off * stride_cache_pos - + kv_head * stride_cache_head + + page_off.to(tl.int64) * stride_cache_pos + + tl.cast(kv_head, tl.int64) * stride_cache_head ) # ============================================================ @@ -356,11 +356,11 @@ def _tq_full_dequant_kv( page_idx = pos // BLOCK_SIZE page_off = pos % BLOCK_SIZE - block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx) + block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx).to(tl.int64) slot_base = ( block_num * stride_cache_block - + page_off * stride_cache_pos - + hid * stride_cache_head + + tl.cast(page_off, tl.int64) * stride_cache_pos + + tl.cast(hid, tl.int64) * stride_cache_head ) d_offs = tl.arange(0, BLOCK_D) diff --git a/vllm/v1/attention/ops/triton_turboquant_store.py b/vllm/v1/attention/ops/triton_turboquant_store.py index 3da3347d5df..3ad2d41488e 100644 --- a/vllm/v1/attention/ops/triton_turboquant_store.py +++ b/vllm/v1/attention/ops/triton_turboquant_store.py @@ -174,10 +174,13 @@ def _tq_fused_store_fp8( slot = tl.load(Slot_mapping_ptr + token_idx) if slot < 0: return - blk = slot // BLOCK_SIZE - off = slot % BLOCK_SIZE + blk = (slot // BLOCK_SIZE).to(tl.int64) + off = (slot % BLOCK_SIZE).to(tl.int64) + head_idx_i64 = tl.cast(head_idx, tl.int64) slot_base = ( - blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head + blk * stride_cache_block + + off * stride_cache_pos + + head_idx_i64 * stride_cache_head ) base = pid * D @@ -259,10 +262,13 @@ def _tq_fused_store_mse( slot = tl.load(Slot_mapping_ptr + token_idx) if slot < 0: return - blk = slot // BLOCK_SIZE - off = slot % BLOCK_SIZE + blk = (slot // BLOCK_SIZE).to(tl.int64) + off = (slot % BLOCK_SIZE).to(tl.int64) + head_idx_i64 = tl.cast(head_idx, tl.int64) slot_base = ( - blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head + blk * stride_cache_block + + off * stride_cache_pos + + head_idx_i64 * stride_cache_head ) base = pid * D From 5cdddddd4a03b0d1b6fa1940458e432b91457ae1 Mon Sep 17 00:00:00 2001 From: Yanan Cao Date: Fri, 17 Apr 2026 14:36:49 -0700 Subject: [PATCH 096/150] [Kernel] [Helion] Force disable HOP path due to performance regression (#40171) Signed-off-by: Yanan Cao Co-authored-by: Claude Sonnet 4 --- vllm/kernels/helion/register.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/kernels/helion/register.py b/vllm/kernels/helion/register.py index 080c4cdda0f..30dcbe08c40 100644 --- a/vllm/kernels/helion/register.py +++ b/vllm/kernels/helion/register.py @@ -53,14 +53,15 @@ if not has_helion(): ) import helion -from helion._compat import requires_torch_version from helion.autotuner.base_search import BaseAutotuner from helion.runtime.config import Config from helion.runtime.settings import default_autotuner_fn # TODO(gmagogsfm): Remove CustomOp fallback path (_get_or_register_custom_op, # vllm_helion_lib, direct_register_custom_op) once vLLM requires PyTorch >= 2.11. -_HOP_AVAILABLE = requires_torch_version("2.11") +# FIXME(gmagogsfm): Re-enable HOP path once performance regression is fixed. +# _HOP_AVAILABLE = requires_torch_version("2.11") +_HOP_AVAILABLE = False if _HOP_AVAILABLE: from helion._compat import supports_torch_compile_fusion From a8bffaa13337236408e7fe6387e1da764191b05a Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Fri, 17 Apr 2026 19:42:32 -0400 Subject: [PATCH 097/150] [Kernel] Add MXFP4 W4A4 CUTLASS MoE kernel for SM100 (#37463) Signed-off-by: mgoin --- .buildkite/test_areas/kernels.yaml | 1 + CMakeLists.txt | 4 +- csrc/libtorch_stable/ops.h | 23 + .../fp4/mxfp4_blockwise_moe_kernel.cu | 468 ++++++++++++++++++ .../quantization/fp4/mxfp4_experts_quant.cu | 422 ++++++++++++++++ csrc/libtorch_stable/torch_bindings.cpp | 22 + tests/kernels/moe/test_mxfp4_moe.py | 248 ++++++++++ vllm/_custom_ops.py | 135 +++++ .../model_executor/layers/fused_moe/config.py | 19 + .../layers/fused_moe/cutlass_moe.py | 295 +++++++++++ .../compressed_tensors_moe_w4a4_mxfp4.py | 78 ++- 11 files changed, 1701 insertions(+), 14 deletions(-) create mode 100644 csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu create mode 100644 csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu create mode 100644 tests/kernels/moe/test_mxfp4_moe.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 3e7376e4134..2df0a5e253f 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -141,6 +141,7 @@ steps: - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py + - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py - pytest -v -s tests/kernels/moe/test_flashinfer.py - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py diff --git a/CMakeLists.txt b/CMakeLists.txt index fd4314bec52..fbc335b7f8e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -952,7 +952,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu") + "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" + "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu" + "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu") set_gencode_flags_for_srcs( SRCS "${SRCS}" CUDA_ARCHS "${FP4_ARCHS}") diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 8153102c598..fdf628d6fd9 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -134,4 +134,27 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, torch::stable::Tensor& input, torch::stable::Tensor& input_global_scale); +void mxfp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts, + int64_t n_experts); + +void silu_and_mul_mxfp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts, + int64_t n_experts); + +void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets); + #endif diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu new file mode 100644 index 00000000000..8a493fdf22c --- /dev/null +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu @@ -0,0 +1,468 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project + * + * MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100. + * Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups. + */ + +#include +#include +#include "libtorch_stable/torch_utils.h" + +#include + +#include "cutlass_extensions/common.hpp" + +#include "cute/tensor.hpp" +#include "cutlass/tensor_ref.h" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/group_array_problem_shape.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" + +#include "cutlass/util/packed_stride.hpp" +#include + +using namespace cute; + +// Offset-computation kernel for MXFP4 grouped GEMM (group size 32). +template +__global__ void __mxfp4_get_group_gemm_starts( + ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets, + ElementSF** a_scales_offsets, ElementSF** b_scales_offsets, + LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int, + ElementAB* a_base_as_int, ElementAB* b_base_as_int, + ElementC* out_base_as_int, ElementSF* a_scales_base_as_int, + ElementSF* b_scales_base_as_int, const int32_t* expert_offsets, + const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes, + int64_t* a_strides, int64_t* b_strides, int64_t* c_strides, + const int64_t a_stride_val, const int64_t b_stride_val, + const int64_t c_stride_val, const int K, const int N) { + int64_t expert_id = threadIdx.x; + if (expert_id >= gridDim.x * blockDim.x) { + return; + } + int64_t expert_offset = static_cast(expert_offsets[expert_id]); + int64_t sf_offset = static_cast(sf_offsets[expert_id]); + int64_t group_size = 32; + int64_t m = static_cast(problem_sizes_as_shapes[expert_id * 3]); + int64_t n = static_cast(problem_sizes_as_shapes[expert_id * 3 + 1]); + int64_t k = static_cast(problem_sizes_as_shapes[expert_id * 3 + 2]); + assert((m >= 0 && n == N && k == K && k % 2 == 0) && + "unexpected problem sizes"); + + int64_t half_k = static_cast(k / 2); + int64_t group_k = static_cast(k / group_size); + // Shape of A as uint8/byte = [M, K // 2] + a_offsets[expert_id] = a_base_as_int + expert_offset * half_k; + // Shape of B as uint8/byte = [E, N, K // 2] + b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k; + // Shape of C = [M, N] + out_offsets[expert_id] = out_base_as_int + expert_offset * n; + // Shape of a_scale = [sum(sf_sizes), K // group_size] + a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k; + + assert((reinterpret_cast(a_scales_offsets[expert_id]) % 128) == + 0 && + "TMA requires 128-byte alignment"); + + // Shape of B scale = [E, N, K // group_size] + b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k; + assert((reinterpret_cast(b_scales_offsets[expert_id]) % 128) == + 0 && + "TMA requires 128-byte alignment"); + + // Initialize strides + a_strides[expert_id] = a_stride_val; + b_strides[expert_id] = b_stride_val; + c_strides[expert_id] = c_stride_val; + + LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id; + LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id; + + *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape( + static_cast(m), static_cast(n), static_cast(k), 1)); + *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape( + static_cast(m), static_cast(n), static_cast(k), 1)); +} + +#define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE, \ + TENSOR_C_TYPE, C_TYPE, LayoutSFA, \ + LayoutSFB, ScaleConfig) \ + else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \ + __mxfp4_get_group_gemm_starts \ + <<<1, num_experts, 0, stream>>>( \ + static_cast(a_starts.data_ptr()), \ + static_cast(b_starts.data_ptr()), \ + static_cast(out_starts.data_ptr()), \ + static_cast(a_scales_starts.data_ptr()), \ + static_cast(b_scales_starts.data_ptr()), \ + reinterpret_cast(layout_sfa.data_ptr()), \ + reinterpret_cast(layout_sfb.data_ptr()), \ + static_cast(a_tensors.data_ptr()), \ + static_cast(b_tensors.data_ptr()), \ + static_cast(out_tensors.data_ptr()), \ + static_cast(a_scales.data_ptr()), \ + static_cast(b_scales.data_ptr()), \ + static_cast(expert_offsets.data_ptr()), \ + static_cast(sf_offsets.data_ptr()), \ + static_cast(problem_sizes.data_ptr()), \ + static_cast(a_strides.data_ptr()), \ + static_cast(b_strides.data_ptr()), \ + static_cast(c_strides.data_ptr()), a_stride_val, \ + b_stride_val, c_stride_val, K, N); \ + } + +template +void mxfp4_run_get_group_gemm_starts( + const torch::stable::Tensor& a_starts, + const torch::stable::Tensor& b_starts, + const torch::stable::Tensor& out_starts, + const torch::stable::Tensor& a_scales_starts, + const torch::stable::Tensor& b_scales_starts, + const torch::stable::Tensor& layout_sfa, + const torch::stable::Tensor& layout_sfb, + const torch::stable::Tensor& a_strides, + const torch::stable::Tensor& b_strides, + const torch::stable::Tensor& c_strides, int64_t a_stride_val, + int64_t b_stride_val, int64_t c_stride_val, + torch::stable::Tensor const& a_tensors, + torch::stable::Tensor const& b_tensors, + torch::stable::Tensor const& out_tensors, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + torch::stable::Tensor const& expert_offsets, + torch::stable::Tensor const& sf_offsets, + torch::stable::Tensor const& problem_sizes, int M, int N, int K) { + int num_experts = (int)expert_offsets.size(0); + auto stream = get_current_cuda_stream(a_tensors.get_device_index()); + + STD_TORCH_CHECK(out_tensors.size(1) == N, + "Output tensor shape doesn't match expected shape"); + STD_TORCH_CHECK(K / 2 == b_tensors.size(2), + "b_tensors(dim = 2) and a_tensors(dim = 1) trailing" + " dimension must match"); + if (false) { + } + // MXFP4 uses E8M0 (float_ue8m0_t) scale factors + __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t, + torch::headeronly::ScalarType::BFloat16, + cutlass::bfloat16_t, LayoutSFA, LayoutSFB, + ScaleConfig) + __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t, + torch::headeronly::ScalarType::Half, half, + LayoutSFA, LayoutSFB, ScaleConfig) + else { + STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)"); + } +} + +template +void run_mxfp4_blockwise_scaled_group_mm_sm100( + torch::stable::Tensor& output, const torch::stable::Tensor& a, + const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets, int M, int N, int K) { + using ProblemShape = + cutlass::gemm::GroupProblemShape>; + using ElementType = cutlass::float_e2m1_t; + using ElementSFType = cutlass::float_ue8m0_t; + using ElementA = cutlass::mx_float4_t; + using ElementB = cutlass::mx_float4_t; + + using ElementC = OutType; + using ElementD = ElementC; + using ElementAccumulator = float; + // Layout definitions + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using LayoutD = LayoutC; + + static constexpr int AlignmentA = 32; + static constexpr int AlignmentB = 32; + static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + + // Architecture definitions + using ArchTag = cutlass::arch::Sm100; + using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp; + using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; + using StageCountType = cutlass::gemm::collective::StageCountAuto; + + using ClusterShape = Shape<_1, _1, _1>; + struct MMA1SMConfig { + using MmaTileShape = Shape<_128, _128, _128>; + using KernelSchedule = + cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; + }; + + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape, + ClusterShape, Shape<_128, _64>, ElementAccumulator, + ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD, + LayoutC*, AlignmentD, + typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA, + ElementB, LayoutB*, AlignmentB, ElementAccumulator, + typename MMA1SMConfig::MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + typename MMA1SMConfig::KernelSchedule>::CollectiveOp; + + using GemmKernel = + cutlass::gemm::kernel::GemmUniversal; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using StrideA = typename Gemm::GemmKernel::InternalStrideA; + using StrideB = typename Gemm::GemmKernel::InternalStrideB; + using StrideC = typename Gemm::GemmKernel::InternalStrideC; + using StrideD = typename Gemm::GemmKernel::InternalStrideD; + + using LayoutSFA = + typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; + using LayoutSFB = + typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; + using ScaleConfig = + typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape; + int num_experts = static_cast(expert_offsets.size(0)); + + torch::stable::Tensor a_ptrs = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor b_ptrs = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor out_ptrs = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor a_scales_ptrs = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor b_scales_ptrs = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor layout_sfa = torch::stable::empty( + {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt, + a.device()); + torch::stable::Tensor layout_sfb = torch::stable::empty( + {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt, + a.device()); + torch::stable::Tensor a_strides1 = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor b_strides1 = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor c_strides1 = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + + mxfp4_run_get_group_gemm_starts( + a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa, + layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2, + b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale, + b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K); + + // Create an instance of the GEMM + Gemm gemm_op; + + UnderlyingProblemShape* problem_sizes_as_shapes = + static_cast(problem_sizes.data_ptr()); + + // Set the Scheduler info + cutlass::KernelHardwareInfo hw_info; + using RasterOrderOptions = typename cutlass::gemm::kernel::detail:: + PersistentTileSchedulerSm100GroupParams< + typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions; + typename Gemm::GemmKernel::TileSchedulerArguments scheduler; + scheduler.raster_order = RasterOrderOptions::AlongM; + hw_info.device_id = a.get_device_index(); + static std::unordered_map cached_sm_counts; + if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) { + cached_sm_counts[hw_info.device_id] = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count( + hw_info.device_id); + } + hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX); + + // Mainloop Arguments + typename GemmKernel::MainloopArguments mainloop_args{ + static_cast(a_ptrs.data_ptr()), + static_cast(a_strides1.data_ptr()), + static_cast(b_ptrs.data_ptr()), + static_cast(b_strides1.data_ptr()), + static_cast(a_scales_ptrs.data_ptr()), + reinterpret_cast(layout_sfa.data_ptr()), + static_cast(b_scales_ptrs.data_ptr()), + reinterpret_cast(layout_sfb.data_ptr())}; + + // Epilogue Arguments + typename GemmKernel::EpilogueArguments epilogue_args{ + {}, // epilogue.thread + nullptr, + static_cast(c_strides1.data_ptr()), + static_cast(out_ptrs.data_ptr()), + static_cast(c_strides1.data_ptr())}; + auto& fusion_args = epilogue_args.thread; + // Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C + fusion_args.alpha_ptr = nullptr; + fusion_args.beta_ptr = nullptr; + fusion_args.alpha = 1.0f; + fusion_args.alpha_ptr_array = nullptr; + fusion_args.dAlpha = {_0{}, _0{}, 0}; + fusion_args.beta = 0.0f; + fusion_args.beta_ptr_array = nullptr; + fusion_args.dBeta = {_0{}, _0{}, 0}; + + // Gemm Arguments + typename GemmKernel::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGrouped, + {num_experts, problem_sizes_as_shapes, nullptr}, + mainloop_args, + epilogue_args, + hw_info, + scheduler}; + + size_t workspace_size = Gemm::get_workspace_size(args); + auto workspace = + torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte, + std::nullopt, a.device()); + const cudaStream_t stream = get_current_cuda_stream(a.get_device_index()); + + auto can_implement_status = gemm_op.can_implement(args); + STD_TORCH_CHECK( + can_implement_status == cutlass::Status::kSuccess, + "Failed to implement MXFP4 GEMM: status=", (int)can_implement_status); + + // Run the GEMM + auto status = gemm_op.initialize(args, workspace.data_ptr()); + STD_TORCH_CHECK(status == cutlass::Status::kSuccess, + "Failed to initialize MXFP4 GEMM: status=", (int)status, + " workspace_size=", workspace_size, + " num_experts=", num_experts, " M=", M, " N=", N, " K=", K); + + status = gemm_op.run(args, workspace.data_ptr(), stream); + STD_TORCH_CHECK(status == cutlass::Status::kSuccess, + "Failed to run MXFP4 GEMM"); +} + +template +void run_mxfp4_blockwise_scaled_group_mm( + torch::stable::Tensor& output, const torch::stable::Tensor& a, + const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets, int M, int N, int K) { + int32_t version_num = get_sm_version_num(); +#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 + if (version_num >= 100 && version_num < 120) { + run_mxfp4_blockwise_scaled_group_mm_sm100( + output, a, b, a_blockscale, b_blockscales, problem_sizes, + expert_offsets, sf_offsets, M, N, K); + return; + } +#endif + STD_TORCH_CHECK_NOT_IMPLEMENTED( + false, + "No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ", + version_num, ". Required capability: 100"); +} + +#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 +constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte; +// E8M0 scale factors stored as uint8 +constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte; +#endif + +#define CHECK_TYPE(x, st, m) \ + STD_TORCH_CHECK(x.scalar_type() == st, \ + ": Inconsistency of torch::stable::Tensor type:", m) +#define CHECK_TH_CUDA(x, m) \ + STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.") +#define CHECK_CONTIGUOUS(x, m) \ + STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.") +#define CHECK_INPUT(x, st, m) \ + CHECK_TH_CUDA(x, m); \ + CHECK_CONTIGUOUS(x, m); \ + CHECK_TYPE(x, st, m) + +void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets) { +#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 + // Input validation + CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a"); + CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b"); + // MXFP4 uses E8M0 scale factors (stored as uint8) + CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale"); + CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales"); + + STD_TORCH_CHECK( + a_blockscale.dim() == 2, + "expected a_blockscale to be of shape [num_experts, rounded_m," + " k // group_size], observed rank: ", + a_blockscale.dim()) + STD_TORCH_CHECK(b_blockscales.dim() == 3, + "expected b_blockscale to be of shape: " + " [num_experts, n, k // group_size], observed rank: ", + b_blockscales.dim()) + STD_TORCH_CHECK(problem_sizes.dim() == 2, + "problem_sizes must be a 2D tensor"); + STD_TORCH_CHECK(problem_sizes.size(1) == 3, + "problem_sizes must have the shape (num_experts, 3)"); + STD_TORCH_CHECK( + problem_sizes.size(0) == expert_offsets.size(0), + "Number of experts in problem_sizes must match expert_offsets"); + STD_TORCH_CHECK( + problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int, + "problem_sizes must be int32."); + + int M = static_cast(a.size(0)); + int N = static_cast(b.size(1)); + int E = static_cast(b.size(0)); + int K = static_cast(2 * b.size(2)); + + if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { + run_mxfp4_blockwise_scaled_group_mm( + output, a, b, a_blockscale, b_blockscales, problem_sizes, + expert_offsets, sf_offsets, M, N, K); + } else { + run_mxfp4_blockwise_scaled_group_mm( + output, a, b, a_blockscale, b_blockscales, problem_sizes, + expert_offsets, sf_offsets, M, N, K); + } +#else + STD_TORCH_CHECK_NOT_IMPLEMENTED( + false, + "No compiled cutlass_mxfp4_group_mm kernel; build vLLM with " + "SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+."); +#endif +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm)); +} diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu new file mode 100644 index 00000000000..9119f28a750 --- /dev/null +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu @@ -0,0 +1,422 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project + * + * MXFP4 activation quantization kernel for MoE experts. + * Quantizes BF16/FP16 activations to MXFP4: E2M1 values with E8M0 block scales + * over 32-element groups. + * + * Uses PACK16 E2M1 conversion helpers (nvfp4_utils.cuh) configured for: + * - Block size 32 (2 threads per SF in PACK16 mode) + * - E8M0 (power-of-two) scale factors + * - SF layout: [numMTiles, numKTiles, 32, 4, 4] where numKTiles=ceil(K/128) + */ + +// MXFP4 requires PACK16 mode (16 elements per thread) so that +// 2 threads cover 32-element blocks. This requires CUDA >= 12.9. +// Must be defined before any header that (transitively) includes +// nvfp4_utils.cuh. +#define NVFP4_ENABLE_ELTS16 1 + +#include +#include +#include +#include + +#include +#include "libtorch_stable/torch_utils.h" +#include "libtorch_stable/dispatch_utils.h" +#include "cuda_vec_utils.cuh" +#include "cuda_utils.h" + +#include "nvfp4_utils.cuh" +static_assert(CVT_FP4_ELTS_PER_THREAD == 16, + "MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)"); + +#include "launch_bounds_utils.h" + +namespace vllm { + +// MXFP4 block size constants +static constexpr int MXFP4_SF_VEC_SIZE = 32; + +// For PACK16 mode (CVT_FP4_ELTS_PER_THREAD=16): 2 threads per SF +// For PACK8 mode (CVT_FP4_ELTS_PER_THREAD=8): 4 threads per SF +static constexpr int MXFP4_NUM_THREADS_PER_SF = + MXFP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD; + +// MXFP4 quantization kernel for experts. +// Uses 32-element blocks with E8M0 (UE8M0) scale factors. +// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses +// SiLU(gate)*up before quantization. +template +__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) + mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in, + fp4_packed_t* out, uint32_t* SFout, + uint32_t* input_offset_by_experts, + uint32_t* output_scale_offset_by_experts, + int n_experts, bool low_latency) { + using PackedVec = PackedVec; + static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, + "Vec size is not matched."); + + // MXFP4: numKTiles = ceil(numCols / 128) since block_size=32, 4 SFs/tile + int32_t const numKTiles = (numCols + 127) / 128; + + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; + int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow; + + for (int globalIdx = tid; globalIdx < numRows * colsPerRow; + globalIdx += gridDim.x * blockDim.x) { + int rowIdx = globalIdx / colsPerRow; + int colIdx = globalIdx % colsPerRow; + + int rowIdx_in_expert = 0; + int expert_idx = 0; + + if constexpr (SMALL_NUM_EXPERTS) { + for (int i = 0; i < n_experts; i++) { + uint32_t current_offset = __ldca(&input_offset_by_experts[i]); + uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]); + if (rowIdx >= current_offset && rowIdx < next_offset) { + rowIdx_in_expert = rowIdx - current_offset; + expert_idx = i; + break; + } + } + } else { + uint32_t local_offsets[17]; + for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) { + *reinterpret_cast(local_offsets) = + __ldca(reinterpret_cast( + &input_offset_by_experts[chunk_start])); + *reinterpret_cast(local_offsets + 4) = + __ldca(reinterpret_cast( + &input_offset_by_experts[chunk_start + 4])); + *reinterpret_cast(local_offsets + 8) = + __ldca(reinterpret_cast( + &input_offset_by_experts[chunk_start + 8])); + *reinterpret_cast(local_offsets + 12) = + __ldca(reinterpret_cast( + &input_offset_by_experts[chunk_start + 12])); + local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]); + +#pragma unroll + for (int i = 0; i < 16; i++) { + if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) { + rowIdx_in_expert = rowIdx - local_offsets[i]; + expert_idx = chunk_start + i; + break; + } + } + } + } + + // Load input and optionally apply fused SiLU+Mul + int64_t inOffset = rowIdx * inColsPerRow + colIdx; + PackedVec in_vec = reinterpret_cast(in)[inOffset]; + PackedVec quant_input; + if constexpr (FUSE_SILU_MUL) { + PackedVec in_vec_up = + reinterpret_cast(in)[inOffset + colsPerRow]; + quant_input = compute_silu_mul(in_vec, in_vec_up); + } else { + quant_input = in_vec; + } + + // In PACK16 mode, each thread outputs 16 E2M1 values = u32x2 + int64_t outOffset = rowIdx * colsPerRow + colIdx; + auto& out_pos = out[outOffset]; + + uint32_t* SFout_in_expert = + SFout + output_scale_offset_by_experts[expert_idx] * numKTiles; + + // Use MXFP4_NUM_THREADS_PER_SF (2 for PACK16) for 32-element blocks + auto sf_out = + cvt_quant_to_fp4_get_sf_out_offset( + rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert); + + // Block E8M0 scales only; no extra tensor-level scale in this path + constexpr float SFScaleVal = 1.0f; + // UE8M0_SF=true for MXFP4 E8M0 scale factors + out_pos = + cvt_warp_fp16_to_fp4( + quant_input, SFScaleVal, sf_out); + } +} + +// Large M_topk variant using shared memory for expert offsets +template +__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024)) + mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in, + fp4_packed_t* out, uint32_t* SFout, + uint32_t* input_offset_by_experts, + uint32_t* output_scale_offset_by_experts, + int n_experts) { + using PackedVec = PackedVec; + static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, + "Vec size is not matched."); + + // MXFP4: numKTiles = ceil(numCols / 128) + int32_t const numKTiles = (numCols + 127) / 128; + + extern __shared__ uint32_t shared_input_offsets[]; + + if constexpr (SMALL_NUM_EXPERTS) { + for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) { + shared_input_offsets[i] = input_offset_by_experts[i]; + } + } else { + for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) { + *reinterpret_cast(&shared_input_offsets[i]) = + *reinterpret_cast(&input_offset_by_experts[i]); + } + if (threadIdx.x == 0) { + shared_input_offsets[n_experts] = input_offset_by_experts[n_experts]; + } + } + + __syncthreads(); + + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; + int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow; + + for (int globalIdx = tid; globalIdx < numRows * colsPerRow; + globalIdx += gridDim.x * blockDim.x) { + int rowIdx = globalIdx / colsPerRow; + int colIdx = globalIdx % colsPerRow; + + int rowIdx_in_expert = 0; + int expert_idx = 0; + + // Binary search through experts using shared memory + int left = 0, right = n_experts - 1; + while (left <= right) { + int mid = (left + right) / 2; + uint32_t mid_offset = shared_input_offsets[mid]; + uint32_t next_offset = shared_input_offsets[mid + 1]; + + if (rowIdx >= mid_offset && rowIdx < next_offset) { + rowIdx_in_expert = rowIdx - mid_offset; + expert_idx = mid; + break; + } else if (rowIdx < mid_offset) { + right = mid - 1; + } else { + left = mid + 1; + } + } + + int64_t inOffset = rowIdx * inColsPerRow + colIdx; + PackedVec in_vec = reinterpret_cast(in)[inOffset]; + PackedVec quant_input; + if constexpr (FUSE_SILU_MUL) { + PackedVec in_vec_up = + reinterpret_cast(in)[inOffset + colsPerRow]; + quant_input = compute_silu_mul(in_vec, in_vec_up); + } else { + quant_input = in_vec; + } + + int64_t outOffset = rowIdx * colsPerRow + colIdx; + auto& out_pos = out[outOffset]; + + // MXFP4 has no global scale - only block-level E8M0 scale factors + constexpr float SFScaleVal = 1.0f; + + uint32_t* SFout_in_expert = + SFout + output_scale_offset_by_experts[expert_idx] * numKTiles; + + auto sf_out = + cvt_quant_to_fp4_get_sf_out_offset( + rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert); + + out_pos = + cvt_warp_fp16_to_fp4( + quant_input, SFScaleVal, sf_out); + } +} + +template +void mxfp4_quant_impl(void* output, void* output_scale, void* input, + void* input_offset_by_experts, + void* output_scale_offset_by_experts, int m_topk, int k, + int n_experts, cudaStream_t stream) { + int multiProcessorCount = + get_device_attribute(cudaDevAttrMultiProcessorCount, -1); + + int const workSizePerRow = k / ELTS_PER_THREAD; + int const totalWorkSize = m_topk * workSizePerRow; + dim3 block(std::min(workSizePerRow, 512)); + int const numBlocksPerSM = + vllm_runtime_blocks_per_sm(static_cast(block.x)); + dim3 grid(std::min(static_cast((totalWorkSize + block.x - 1) / block.x), + multiProcessorCount * numBlocksPerSM)); + while (grid.x <= multiProcessorCount && block.x > 64) { + grid.x *= 2; + block.x = (block.x + 1) / 2; + } + + int const blockRepeat = + (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x); + if (blockRepeat > 1) { + size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t); + if (n_experts >= 4) { + mxfp4_cvt_fp16_to_fp4 + <<>>( + m_topk, k, reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(output_scale), + reinterpret_cast(input_offset_by_experts), + reinterpret_cast(output_scale_offset_by_experts), + n_experts); + } else { + mxfp4_cvt_fp16_to_fp4 + <<>>( + m_topk, k, reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(output_scale), + reinterpret_cast(input_offset_by_experts), + reinterpret_cast(output_scale_offset_by_experts), + n_experts); + } + } else { + if (n_experts >= 16) { + mxfp4_cvt_fp16_to_fp4 + <<>>( + m_topk, k, reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(output_scale), + reinterpret_cast(input_offset_by_experts), + reinterpret_cast(output_scale_offset_by_experts), + n_experts, /* bool low_latency */ true); + } else { + mxfp4_cvt_fp16_to_fp4<<>>( + m_topk, k, reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(output_scale), + reinterpret_cast(input_offset_by_experts), + reinterpret_cast(output_scale_offset_by_experts), + n_experts, /* bool low_latency */ true); + } + } +} + +} // namespace vllm + +/*Quantization entry for mxfp4 experts quantization*/ +#define CHECK_TH_CUDA(x, m) \ + STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x, m) \ + STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous") +#define CHECK_INPUT(x, m) \ + CHECK_TH_CUDA(x, m); \ + CHECK_CONTIGUOUS(x, m); + +constexpr auto HALF = torch::headeronly::ScalarType::Half; +constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16; +constexpr auto INT = torch::headeronly::ScalarType::Int; +constexpr auto UINT8 = torch::headeronly::ScalarType::Byte; + +static constexpr int MXFP4_BLOCK_SIZE = 32; + +static void validate_mxfp4_experts_quant_inputs( + torch::stable::Tensor const& output, + torch::stable::Tensor const& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts, + int64_t n_experts, int64_t m_topk, int64_t k) { + CHECK_INPUT(output, "output"); + CHECK_INPUT(output_scale, "output_scale"); + CHECK_INPUT(input, "input"); + CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts"); + CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts"); + + STD_TORCH_CHECK(output.dim() == 2); + STD_TORCH_CHECK(output_scale.dim() == 2); + STD_TORCH_CHECK(input.dim() == 2); + STD_TORCH_CHECK(input_offset_by_experts.dim() == 1); + STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1); + + STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16); + STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT); + STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT); + // output is uint8 (two mxfp4 values packed into one uint8) + // output_scale is int32 (four E8M0 values packed into one int32) + STD_TORCH_CHECK(output.scalar_type() == UINT8); + STD_TORCH_CHECK(output_scale.scalar_type() == INT); + + STD_TORCH_CHECK(k % MXFP4_BLOCK_SIZE == 0, "k must be a multiple of 32"); + STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1); + STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1); + STD_TORCH_CHECK(output.size(0) == m_topk); + STD_TORCH_CHECK(output.size(1) == k / 2); + int scales_k = k / MXFP4_BLOCK_SIZE; + // K-dimension scale columns padded to a multiple of 4 for swizzle layout + int padded_k = (scales_k + (4 - 1)) / 4 * 4; + // 4 = 4 E8M0 values packed into one int32 + STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k); +} + +void mxfp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts, + int64_t n_experts) { + auto m_topk = input.size(0); + auto k = input.size(1); + + validate_mxfp4_experts_quant_inputs( + output, output_scale, input, input_offset_by_experts, + output_scale_offset_by_experts, n_experts, m_topk, k); + + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(input.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + input.scalar_type(), "mxfp4_experts_quant_kernel", [&] { + using cuda_type = vllm::CUDATypeConverter::Type; + vllm::mxfp4_quant_impl( + output.data_ptr(), output_scale.data_ptr(), input.data_ptr(), + input_offset_by_experts.data_ptr(), + output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts, + stream); + }); +} + +void silu_and_mul_mxfp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts, + int64_t n_experts) { + auto m_topk = input.size(0); + auto k_times_2 = input.size(1); + STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)"); + auto k = k_times_2 / 2; + + validate_mxfp4_experts_quant_inputs( + output, output_scale, input, input_offset_by_experts, + output_scale_offset_by_experts, n_experts, m_topk, k); + + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(input.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + input.scalar_type(), "silu_mul_mxfp4_experts_quant_kernel", [&] { + using cuda_type = vllm::CUDATypeConverter::Type; + vllm::mxfp4_quant_impl( + output.data_ptr(), output_scale.data_ptr(), input.data_ptr(), + input_offset_by_experts.data_ptr(), + output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts, + stream); + }); +} diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index c31844948e5..95ed6b44f10 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -116,6 +116,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas," " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE) + ops.def( + "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b," + " Tensor a_blockscale, Tensor b_blockscales," + " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + // Compute NVFP4 block quantized tensor. ops.def( "scaled_fp4_quant(Tensor input," @@ -149,6 +155,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," "Tensor output_scale_offset_by_experts) -> ()"); + // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs). + ops.def( + "mxfp4_experts_quant(Tensor! output, Tensor! output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + + // Fused SiLU+Mul+MXFP4 experts quantization. + ops.def( + "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! " + "output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + // Fused SiLU+Mul+NVFP4 quantization. ops.def( "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, " @@ -233,6 +252,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("silu_and_mul_scaled_fp4_experts_quant", TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant)); ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant)); + ops.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant)); + ops.impl("silu_and_mul_mxfp4_experts_quant", + TORCH_BOX(&silu_and_mul_mxfp4_experts_quant)); // W4A8 ops: impl registrations are in the source files // (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu) diff --git a/tests/kernels/moe/test_mxfp4_moe.py b/tests/kernels/moe/test_mxfp4_moe.py new file mode 100644 index 00000000000..11fd853f54f --- /dev/null +++ b/tests/kernels/moe/test_mxfp4_moe.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Tests for SM100 CUTLASS MXFP4 x MXFP4 grouped MoE kernels.""" + +import random + +import pytest +import torch + +from tests.kernels.utils import torch_moe_single +from vllm import _custom_ops as ops +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +random.seed(42) +set_random_seed(42) + +MXFP4_BLOCK_SIZE = 32 + + +def align(val: int, alignment: int = 128) -> int: + return int((val + alignment - 1) // alignment * alignment) + + +def calc_diff(x, y): + x, y = x.double(), y.double() + denominator = (x * x + y * y).sum() + sim = 2 * (x * y).sum() / denominator + return 1 - sim + + +def is_sm100_supported() -> bool: + return current_platform.is_cuda() and current_platform.is_device_capability_family( + 100 + ) + + +def compute_ref_output( + input_tensor: torch.Tensor, + weight_list: list[torch.Tensor], + expert_offsets: list[int], + expert_offset: int, + num_experts: int, +) -> torch.Tensor: + """Reference output using torch_moe_single with top-1 routing.""" + score = torch.full( + (expert_offset, num_experts), + -1e9, + device=input_tensor.device, + dtype=torch.float32, + ) + for g in range(num_experts): + start = expert_offsets[g] + end = expert_offsets[g + 1] if g + 1 < num_experts else expert_offset + score[start:end, g] = 0.0 + + return torch_moe_single( + input_tensor, torch.stack(weight_list, dim=0), score, topk=1 + ) + + +@pytest.mark.skipif( + not is_sm100_supported(), + reason="cutlass_mxfp4_group_mm requires CUDA SM100", +) +@pytest.mark.parametrize("num_experts", [8, 16, 32]) +@pytest.mark.parametrize("out_dtype", [torch.bfloat16]) +def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype): + """ + Test the MXFP4 grouped GEMM kernel by: + 1. Creating random per-expert inputs and weights + 2. Quantizing both to MXFP4 using the CUDA kernel + 3. Running the CUTLASS grouped GEMM + 4. Comparing against BF16 reference + """ + device = "cuda" + alignment = 128 + # N and K must be multiples of 128 for clean swizzle layout + n_g = random.randint(1, 16) * alignment + k_g = random.randint(1, 16) * alignment + + expert_offset = 0 + expert_offsets_input = [] + problem_sizes = [] + input_list = [] + weight_list = [] + + for g in range(num_experts): + m_g = random.randint(1, 256) + expert_offsets_input.append(expert_offset) + expert_offset += m_g + problem_sizes.append([m_g, n_g, k_g]) + + input_list.append( + torch.normal(0.0, std=0.5, size=(m_g, k_g), device=device, dtype=out_dtype) + ) + weight_list.append( + torch.normal(0.0, std=0.5, size=(n_g, k_g), device=device, dtype=out_dtype) + ) + + input_tensor = torch.concat(input_list, dim=0) # [M_total, K] + + # --- Quantize INPUTS via mxfp4_experts_quant --- + input_bs_offsets = [] + tot = 0 + for g in range(num_experts): + input_bs_offsets.append(tot) + tot += align(problem_sizes[g][0], 128) + input_bs_offsets.append(tot) + + _inp_expert_offsets = torch.tensor( + expert_offsets_input + [expert_offset], device=device, dtype=torch.int32 + ) + _inp_bs_offsets = torch.tensor(input_bs_offsets, device=device, dtype=torch.int32) + + input_quant, input_sf = ops.mxfp4_experts_quant( + input_tensor, + _inp_expert_offsets, + _inp_bs_offsets, + num_experts, + topk=1, + ) + + # --- Quantize WEIGHTS via mxfp4_experts_quant --- + # Treat each expert's N weight rows as an "expert" with N tokens + weight_tensor = torch.concat(weight_list, dim=0) # [E*N, K] + weight_expert_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g] + # N is always multiple of 128, so blockscale offsets are clean + weight_bs_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g] + + _wt_expert_offsets = torch.tensor( + weight_expert_offsets, device=device, dtype=torch.int32 + ) + _wt_bs_offsets = torch.tensor(weight_bs_offsets, device=device, dtype=torch.int32) + + weight_quant, weight_sf = ops.mxfp4_experts_quant( + weight_tensor, + _wt_expert_offsets, + _wt_bs_offsets, + num_experts, + topk=1, + ) + + # Reshape weight quantized data to [E, N, K//2] + weight_quant = weight_quant[: num_experts * n_g].view(num_experts, n_g, k_g // 2) + + # Reshape weight scale factors to [E, N, K//32] + # The quant kernel produces uint8 SF buffer. Each row has K//32 SFs. + scales_per_row = k_g // MXFP4_BLOCK_SIZE + weight_sf_flat = weight_sf.view(-1)[: num_experts * n_g * scales_per_row] + weight_sf_3d = weight_sf_flat.view(num_experts, n_g, scales_per_row) + + # Output + output = torch.empty((expert_offset, n_g), device=device, dtype=out_dtype) + + _problem_sizes = torch.tensor(problem_sizes, device=device, dtype=torch.int32) + _expert_offsets = torch.tensor( + expert_offsets_input, device=device, dtype=torch.int32 + ) + _input_bs = torch.tensor(input_bs_offsets[:-1], device=device, dtype=torch.int32) + + # Run the MXFP4 grouped GEMM + ops.cutlass_mxfp4_moe_mm( + output, + input_quant, + weight_quant, + input_sf, + weight_sf_3d, + _problem_sizes, + _expert_offsets, + _input_bs, + ) + + # Reference: BF16 matmul + ref_output = compute_ref_output( + input_tensor=input_tensor, + weight_list=weight_list, + expert_offsets=expert_offsets_input, + expert_offset=expert_offset, + num_experts=num_experts, + ) + + # Compare per-expert + for g in range(num_experts): + start = expert_offsets_input[g] + end = expert_offsets_input[g + 1] if g + 1 < num_experts else expert_offset + if start == end: + continue + baseline = ref_output[start:end] + actual = output[start:end] + diff = calc_diff(actual, baseline) + print( + f"m_g={end - start} n_g={n_g} k_g={k_g} " + f"num_experts={num_experts}, " + f"out_dtype={out_dtype}, diff={diff:.5f}" + ) + # FP4 quantization is very lossy (~4 bits precision) + # Comparing quantized vs full-precision gives cosine diff of 0.05-0.15 + assert diff < 0.15, f"Expert {g}: diff={diff:.5f} exceeds threshold" + + +@pytest.mark.skipif( + not is_sm100_supported(), + reason="mxfp4_experts_quant requires CUDA SM100", +) +def test_mxfp4_experts_quant_basic(): + """ + Basic smoke test for the MXFP4 experts quantization kernel. + """ + device = "cuda" + num_experts = 4 + k = 256 + tokens_per_expert = 16 + + total_tokens = tokens_per_expert * num_experts + input_tensor = torch.randn(total_tokens, k, device=device, dtype=torch.bfloat16) / 5 + + expert_offsets = [i * tokens_per_expert for i in range(num_experts + 1)] + blockscale_offsets = [ + align(i * tokens_per_expert, 128) for i in range(num_experts + 1) + ] + + _expert_offsets = torch.tensor(expert_offsets, device=device, dtype=torch.int32) + _blockscale_offsets = torch.tensor( + blockscale_offsets, device=device, dtype=torch.int32 + ) + + output, output_sf = ops.mxfp4_experts_quant( + input_tensor, + _expert_offsets, + _blockscale_offsets, + num_experts, + topk=1, + ) + + assert output.shape == (total_tokens, k // 2) + assert output.dtype == torch.uint8 + assert output_sf.dtype == torch.uint8 + assert output.any(), "Quantized output is all zeros" + print( + f"MXFP4 experts quant: output shape={output.shape}, sf shape={output_sf.shape}" + ) + print("PASSED") + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index facac5a192d..a7b6a7059b0 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -1150,6 +1150,38 @@ def cutlass_fp4_moe_mm( ) +def cutlass_mxfp4_moe_mm( + out_tensors: torch.Tensor, + a_tensors: torch.Tensor, + b_tensors: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + problem_sizes: torch.Tensor, + expert_offsets: torch.Tensor, + sf_offsets: torch.Tensor, +): + """ + An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4). + + Uses mx_float4_t types with E8M0 scale factors and 32-element blocks. + - a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte) + - a_/b_scales: E8M0 blockscales (uint8, stored in swizzled layout) + - Epilogue uses scalar alpha=1, beta=0 inside the CUDA op (no global scales). + - expert_offsets/sf_offsets: expert boundary indices + - problem_sizes: (num_experts, 3) with (M, N, K) per expert + """ + return torch.ops._C.cutlass_mxfp4_group_mm( + out_tensors, + a_tensors, + b_tensors, + a_scales, + b_scales, + problem_sizes, + expert_offsets, + sf_offsets, + ) + + def mxfp8_experts_quant( input_tensor: torch.Tensor, problem_sizes: torch.Tensor, @@ -1848,6 +1880,109 @@ def silu_and_mul_scaled_fp4_experts_quant( return output, output_scales +def mxfp4_experts_quant( + input_tensor: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + n_experts: int, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Quantize input tensor to MXFP4 for packed MoE inputs. + Uses 32-element blocks with E8M0 (power-of-two) scale factors. + MXFP4 has no global scale - only block-level E8M0 scale factors. + + Args: + input_tensor: [m_topk, k] BF16/FP16 activations + expert_offsets: [n_experts+1] token boundaries per expert + blockscale_offsets: [n_experts+1] SF row boundaries per expert + n_experts: number of experts + topk: number of top-k experts + Returns: + output: [m_topk, k//2] packed E2M1 values (uint8) + output_scales: E8M0 blockscales in swizzled layout (uint8 view) + """ + assert not current_platform.is_rocm() + assert input_tensor.ndim == 2 + + MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE + m_numtopk, k = input_tensor.shape + + assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, ( + f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT(" + f"{MAX_TOKENS_PER_EXPERT})" + f" for cutlass_moe_mxfp4, observed m_numtopk = {m_numtopk}. Use" + f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value." + ) + scales_k = k // 32 + padded_k = (scales_k + (4 - 1)) // 4 + + output = torch.empty( + m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8 + ) + output_scales = torch.empty( + MAX_TOKENS_PER_EXPERT * topk, + padded_k, + dtype=torch.int32, + device=input_tensor.device, + ) + torch.ops._C.mxfp4_experts_quant( + output, + output_scales, + input_tensor, + expert_offsets, + blockscale_offsets, + n_experts, + ) + # E8M0 SFs are stored as uint8 + output_scales = output_scales.view(torch.uint8) + return output, output_scales + + +def silu_and_mul_mxfp4_experts_quant( + input_tensor: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + n_experts: int, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations. + MXFP4 has no global scale - only block-level E8M0 scale factors. + """ + assert not current_platform.is_rocm() + assert input_tensor.ndim == 2 + + MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE + m_numtopk, k_times_2 = input_tensor.shape + assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)" + k = k_times_2 // 2 + + assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk + scales_k = k // 32 + padded_k = (scales_k + (4 - 1)) // 4 + + output = torch.empty( + m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8 + ) + output_scales = torch.empty( + MAX_TOKENS_PER_EXPERT * topk, + padded_k, + dtype=torch.int32, + device=input_tensor.device, + ) + torch.ops._C.silu_and_mul_mxfp4_experts_quant( + output, + output_scales, + input_tensor, + expert_offsets, + blockscale_offsets, + n_experts, + ) + output_scales = output_scales.view(torch.uint8) + return output, output_scales + + # fp8 def scaled_fp8_quant( input: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index a3b941dfa45..1f077ab80be 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -762,6 +762,25 @@ def nvfp4_moe_quant_config( ) +def mxfp4_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for MXFP4 x MXFP4 MoE. + MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no + separate alphas / global activation scales in this config. + """ + return FusedMoEQuantConfig.make( + "mxfp4", + w1_scale=w1_scale, + w2_scale=w2_scale, + per_act_token_quant=False, + per_out_ch_quant=False, + block_shape=None, + ) + + def nvfp4_w4a16_moe_quant_config( g1_alphas: torch.Tensor, g2_alphas: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index 43082b3675a..fdd802e7da3 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -36,6 +36,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8DynamicTokenSym, kFp8StaticChannelSym, kFp8StaticTensorSym, + kMxfp4Dynamic, + kMxfp4Static, kNvfp4Dynamic, kNvfp4Static, ) @@ -795,6 +797,299 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular): ) +def run_cutlass_moe_mxfp4( + output: torch.Tensor, + a: torch.Tensor, + w1_fp4: torch.Tensor, + w1_blockscale: torch.Tensor, + w2_fp4: torch.Tensor, + w2_blockscale: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + m: int, + n: int, + k: int, + e: int, + device: torch.device, + apply_router_weight_on_input: bool = False, +) -> None: + """MXFP4 x MXFP4 MoE implementation using CUTLASS grouped GEMM.""" + is_gated = activation.is_gated + w1_n = n * 2 if is_gated else n + + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch" + assert w1_fp4.dtype == torch.uint8, "weight 1 must be uint8" + assert w2_fp4.dtype == torch.uint8, "weight 2 must be uint8" + assert ( + w1_fp4.ndim == 3 + and w2_fp4.ndim == 3 + and w1_blockscale.ndim == 3 + and w2_blockscale.ndim == 3 + ), "All Weights must be of rank 3 for cutlass_moe_mxfp4" + m_a, k_a = a.shape + e_w1, w1_n_actual, half_k_w1 = w1_fp4.shape + e_w2, k_w2, half_n_w2 = w2_fp4.shape + + assert e_w1 == e_w2 and e_w1 == e + assert k_a == half_k_w1 * 2 and k == k_w2 + assert w1_n_actual == w1_n and half_n_w2 * 2 == n + assert m == m_a + assert 2 * half_k_w1 == k_w2 + assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype" + assert topk_weights.size(0) == m and topk_ids.size(0) == m + + topk = topk_ids.size(1) + out_dtype = a.dtype + num_topk = topk_ids.size(1) + + expert_offsets = torch.empty((e + 1), dtype=torch.int32, device=device) + blockscale_offsets = torch.empty((e + 1), dtype=torch.int32, device=device) + problem_sizes1 = torch.empty((e, 3), dtype=torch.int32, device=device) + problem_sizes2 = torch.empty((e, 3), dtype=torch.int32, device=device) + + a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) + c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) + + if apply_router_weight_on_input: + assert num_topk == 1, ( + "apply_router_weight_on_input is only implemented for topk=1" + ) + a.mul_(topk_weights.to(out_dtype)) + + ops.get_cutlass_moe_mm_data( + topk_ids, + expert_offsets, + problem_sizes1, + problem_sizes2, + a_map, + c_map, + e, + n, + k, + blockscale_offsets, + is_gated=is_gated, + ) + + a = ops.shuffle_rows(a, a_map) + rep_a_fp4, rep_a_blockscale = ops.mxfp4_experts_quant( + a, + expert_offsets, + blockscale_offsets, + e, + num_topk, + ) + c1 = _resize_cache(workspace13, (m * topk, w1_n)) + c2 = _resize_cache(workspace2, (m * topk, n)) + c3 = _resize_cache(workspace13, (m * topk, k)) + + ops.cutlass_mxfp4_moe_mm( + c1, + rep_a_fp4, + w1_fp4, + rep_a_blockscale, + w1_blockscale, + problem_sizes1, + expert_offsets[:-1], + blockscale_offsets[:-1], + ) + del rep_a_fp4, rep_a_blockscale + if activation == MoEActivation.SILU: + int_fp4, int_blockscale = ops.silu_and_mul_mxfp4_experts_quant( + c1, expert_offsets, blockscale_offsets, e, num_topk + ) + else: + apply_moe_activation(activation, c2, c1) + int_fp4, int_blockscale = ops.mxfp4_experts_quant( + c2, expert_offsets, blockscale_offsets, e, num_topk + ) + + ops.cutlass_mxfp4_moe_mm( + c3, + int_fp4, + w2_fp4, + int_blockscale, + w2_blockscale, + problem_sizes2, + expert_offsets[:-1], + blockscale_offsets[:-1], + ) + del int_fp4, int_blockscale + + c3 = ops.shuffle_rows(c3, c_map) + + assert output.dtype == out_dtype + if not apply_router_weight_on_input: + output.copy_( + ( + c3.view(m, num_topk, k) + * topk_weights.view(m, num_topk, 1).to(out_dtype) + ).sum(dim=1), + non_blocking=True, + ) + else: + output.copy_(c3.view(m, num_topk, k).sum(dim=1), non_blocking=True) + return + + +def swizzle_mxfp4_scales( + scales: torch.Tensor, + N: int, + K: int, +) -> torch.Tensor: + """Swizzle flat [N, K//32] E8M0 scales to CUTLASS tiled layout. + + CUTLASS expects MX scale factors in a tiled layout: + [numMTiles, numKTiles, 32, 4, 4] + where numMTiles = ceil(N/128), numKTiles = ceil(K/128), + and the inner dimensions correspond to the swizzle pattern: + mTileIdx = mIdx / 128 + outerMIdx = mIdx % 32 + innerMIdx = (mIdx / 32) % 4 + kTileIdx = kIdx / 4 + innerKIdx = kIdx % 4 + with kIdx = col_in_scale_space (i.e., index into K//32). + """ + assert scales.dtype == torch.uint8 + num_scale_cols = K // 32 # number of E8M0 scale values per row + + num_m_tiles = (N + 127) // 128 + num_k_tiles = (num_scale_cols + 3) // 4 + + # Pad N to multiple of 128 and scale_cols to multiple of 4 + padded_N = num_m_tiles * 128 + padded_scale_cols = num_k_tiles * 4 + + # Start with flat scales, pad if needed + padded = torch.zeros( + padded_N, padded_scale_cols, dtype=torch.uint8, device=scales.device + ) + padded[:N, :num_scale_cols] = scales + + # Reshape to tile structure: + # [numMTiles, 4, 32, numKTiles, 4] + # mTileIdx, innerMIdx, outerMIdx, kTileIdx, innerKIdx + tiled = padded.reshape(num_m_tiles, 4, 32, num_k_tiles, 4) + # Permute to [numMTiles, numKTiles, 32, 4, 4] + # (outerMIdx, innerMIdx, innerKIdx) + tiled = tiled.permute(0, 3, 2, 1, 4).contiguous() + return tiled.reshape(-1) + + +class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular): + """CUTLASS MXFP4 x MXFP4 fused MoE expert implementation.""" + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def _supports_current_device() -> bool: + p = current_platform + return p.is_cuda() and p.is_device_capability_family(100) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == (kMxfp4Static, kMxfp4Dynamic) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in [ + MoEActivation.SILU, + MoEActivation.GELU, + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUSTEP, + MoEActivation.SILU_NO_MUL, + MoEActivation.GELU_NO_MUL, + MoEActivation.RELU2_NO_MUL, + ] + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return moe_parallel_config.ep_size == 1 + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def supports_expert_map(self) -> bool: + return False + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: + return act_dtype + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + workspace1 = (M * topk, max(2 * N, K)) + workspace2 = (M * topk, N) + output = (M, K) + return (workspace1, workspace2, output) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor | None, + workspace2: torch.Tensor | None, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + e, m, n, k, _ = self.moe_problem_size(hidden_states, w1, w2, topk_ids) + n = w2.shape[2] * 2 + + run_cutlass_moe_mxfp4( + output=output, + a=hidden_states, + w1_fp4=w1, + w1_blockscale=self.w1_scale, + w2_fp4=w2, + w2_blockscale=self.w2_scale, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + workspace13=workspace13, + workspace2=workspace2, + m=m, + n=n, + k=k, + e=e, + device=hidden_states.device, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + + # W4A8 def run_cutlass_moe_w4a8_fp8( output: torch.Tensor, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py index 8cc6d17bcc1..57ebb961d48 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py @@ -4,6 +4,7 @@ import torch +import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( FusedMoE, @@ -11,6 +12,10 @@ from vllm.model_executor.layers.fused_moe import ( ) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, + mxfp4_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.cutlass_moe import ( + CutlassExpertsMxfp4, ) from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( MarlinExperts, @@ -36,7 +41,14 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): super().__init__(moe) self.group_size = 32 self.mxfp4_backend = Mxfp4MoeBackend.MARLIN - self.experts_cls = MarlinExperts + self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device() + self.experts_cls: type[mk.FusedMoEExperts] + if self.use_cutlass_mxfp4: + logger.info_once("Using CutlassExpertsMxfp4 for MXFP4 MoE", scope="local") + self.experts_cls = CutlassExpertsMxfp4 + else: + logger.info_once("Using MarlinExperts for MXFP4 MoE", scope="local") + self.experts_cls = MarlinExperts def create_weights( self, @@ -109,11 +121,19 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: - return make_mxfp4_moe_quant_config( - mxfp4_backend=self.mxfp4_backend, - w1_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - ) + if self.use_cutlass_mxfp4: + # W4A4: both weights and activations quantized to MXFP4 + return mxfp4_moe_quant_config( + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + ) + else: + # W4A16: weight-only via Marlin + return make_mxfp4_moe_quant_config( + mxfp4_backend=self.mxfp4_backend, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + ) def process_weights_after_loading(self, layer: FusedMoE) -> None: layer.w13_weight = torch.nn.Parameter( @@ -126,13 +146,45 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): ) delattr(layer, "w2_weight_packed") - logger.warning_once( - "Your GPU does not have native support for FP4 computation but " - "FP4 quantization is being used. Weight-only FP4 compression " - "will be used leveraging the Marlin kernel. This may degrade " - "performance for compute-heavy workloads." - ) - prepare_moe_fp4_layer_for_marlin(layer) + if self.use_cutlass_mxfp4: + # Swizzle weight scales from flat checkpoint layout [E, N, K//32] + # to CUTLASS tiled layout [E, numMTiles*numKTiles*512]. + from vllm.model_executor.layers.fused_moe.cutlass_moe import ( + swizzle_mxfp4_scales, + ) + + E = layer.w13_weight_scale.shape[0] + w13_N = layer.w13_weight_scale.shape[1] + w13_scale_K = layer.w13_weight_scale.shape[2] + w13_K = w13_scale_K * 32 + + w2_M = layer.w2_weight_scale.shape[1] + w2_scale_N = layer.w2_weight_scale.shape[2] + w2_N = w2_scale_N * 32 + + swizzled_w13 = [] + swizzled_w2 = [] + for e_idx in range(E): + s13 = layer.w13_weight_scale[e_idx] + sw13 = swizzle_mxfp4_scales(s13, w13_N, w13_K) + swizzled_w13.append(sw13.reshape(w13_N, w13_scale_K)) + s2 = layer.w2_weight_scale[e_idx] + sw2 = swizzle_mxfp4_scales(s2, w2_M, w2_N) + swizzled_w2.append(sw2.reshape(w2_M, w2_scale_N)) + layer.w13_weight_scale = torch.nn.Parameter( + torch.stack(swizzled_w13), requires_grad=False + ) + layer.w2_weight_scale = torch.nn.Parameter( + torch.stack(swizzled_w2), requires_grad=False + ) + else: + logger.warning_once( + "Your GPU does not have native support for FP4 computation " + "but FP4 quantization is being used. Weight-only FP4 " + "compression will be used leveraging the Marlin kernel. " + "This may degrade performance for compute-heavy workloads." + ) + prepare_moe_fp4_layer_for_marlin(layer) self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config is not None: From 1f45e83756b99aa730c927102ca4708ffbaba2aa Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Fri, 17 Apr 2026 19:49:43 -0400 Subject: [PATCH 098/150] Remove outdated tests test_mixtral_moe and test_duplicated_ignored_sequence_group (#40175) Signed-off-by: mgoin --- tests/kernels/moe/test_moe.py | 130 ---------------------------------- tests/test_regression.py | 16 ----- 2 files changed, 146 deletions(-) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 35b21320f82..5cce699dd2c 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -14,8 +14,6 @@ import pytest import torch from torch.nn import Parameter from torch.nn import functional as F -from transformers import MixtralConfig -from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock import vllm.model_executor.layers.fused_moe # noqa from tests.kernels.moe.utils import ( @@ -24,10 +22,7 @@ from tests.kernels.moe.utils import ( modular_triton_fused_moe, ) from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe -from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig, set_current_vllm_config -from vllm.distributed.parallel_state import init_distributed_environment -from vllm.forward_context import get_forward_context, set_forward_context from vllm.model_executor.layers.fused_moe import ( MoEActivation, fused_topk, @@ -56,12 +51,10 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import ( marlin_quantize, ) from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights -from vllm.model_executor.models.mixtral import MixtralMoE from vllm.platforms import current_platform from vllm.scalar_type import ScalarType, scalar_types from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed -from vllm.v1.worker.workspace import init_workspace_manager def iterative_moe( @@ -681,129 +674,6 @@ def test_fused_moe_wn16( torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("padding", [True, False]) -@pytest.mark.parametrize( - "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False] -) -@torch.inference_mode() -def test_mixtral_moe( - default_vllm_config, - dist_init, - dtype: torch.dtype, - padding: bool, - use_rocm_aiter: bool, - monkeypatch, -): - """Make sure our Mixtral MoE implementation agrees with the one from - huggingface.""" - - # Explicitly set AITER env var based on test parameter to ensure - # consistent behavior regardless of external environment - monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0") - rocm_aiter_ops.refresh_env_variables() - - if use_rocm_aiter and dtype == torch.float32: - pytest.skip("AITER ROCm test skip for float32") - - monkeypatch.setenv("RANK", "0") - monkeypatch.setenv("LOCAL_RANK", "0") - monkeypatch.setenv("WORLD_SIZE", "1") - monkeypatch.setenv("MASTER_ADDR", "localhost") - monkeypatch.setenv("MASTER_PORT", "12345") - init_distributed_environment() - init_workspace_manager(torch.accelerator.current_device_index()) - - # Instantiate our and huggingface's MoE blocks - vllm_config.compilation_config.static_forward_context = dict() - with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): - config = MixtralConfig() - hf_moe = MixtralSparseMoeBlock(config).to(dtype).to("cuda") - vllm_moe = MixtralMoE( - num_experts=config.num_local_experts, - top_k=config.num_experts_per_tok, - hidden_size=config.hidden_size, - intermediate_size=config.intermediate_size, - params_dtype=dtype, - tp_size=1, - dp_size=1, - ).cuda() - - # Load the weights - vllm_moe.gate.weight.data[:] = hf_moe.gate.weight.data - if isinstance(hf_moe.experts, torch.nn.ModuleList): - # Transformers v4 - for i in range(config.num_local_experts): - weights = ( - hf_moe.experts[i].w1.weight.data, - hf_moe.experts[i].w3.weight.data, - ) - vllm_moe.experts.w13_weight[i][:] = torch.cat(weights, dim=0) - vllm_moe.experts.w2_weight[i][:] = hf_moe.experts[i].w2.weight.data - else: - # Transformers v5 - vllm_moe.experts.w13_weight.data[:] = hf_moe.experts.gate_up_proj.data - vllm_moe.experts.w2_weight.data[:] = hf_moe.experts.down_proj.data - # TODO: remove this line after https://github.com/huggingface/transformers/pull/43622 - hf_moe.experts.config._experts_implementation = "eager" - - # Generate input batch of dimensions [batch_size, seq_len, hidden_dim] - hf_inputs = torch.randn((1, 64, config.hidden_size)).to(dtype).to("cuda") - # vLLM uses 1D query [num_tokens, hidden_dim] - vllm_inputs = hf_inputs.flatten(0, 1) - - # Pad the weight if moe padding is enabled - if padding: - vllm_moe.experts.w13_weight = Parameter( - F.pad(vllm_moe.experts.w13_weight, (0, 128), "constant", 0)[ - ..., 0:-128 - ], - requires_grad=False, - ) - vllm_moe.experts.w2_weight = Parameter( - F.pad(vllm_moe.experts.w2_weight, (0, 128), "constant", 0)[..., 0:-128], - requires_grad=False, - ) - torch.accelerator.synchronize() - torch.accelerator.empty_cache() - - # FIXME (zyongye) fix this after we move self.kernel - # assignment in FusedMoE.__init__ - - vllm_moe.experts.quant_method.process_weights_after_loading(vllm_moe.experts) - - # need to override the forward context for unittests, otherwise it assumes - # we're running the model forward pass (the model specified in vllm_config) - get_forward_context().all_moe_layers = None - - # Run forward passes for both MoE blocks - hf_states = hf_moe.forward(hf_inputs) - if isinstance(hf_states, tuple): - # Transformers v4 - hf_states = hf_states[0] - vllm_states = vllm_moe.forward(vllm_inputs) - - mixtral_moe_tol = { - torch.float32: 1e-3, - torch.float16: 1e-3, - torch.bfloat16: 1e-2, - } - - if use_rocm_aiter: - # The values of rtol and atol are set based on the tests in ROCM AITER package. - # https://github.com/ROCm/aiter/blob/dfed377f4be7da96ca2d75ac0761f569676f7240/op_tests/test_moe.py#L174 - torch.testing.assert_close( - hf_states.flatten(0, 1), vllm_states, rtol=0.01, atol=100 - ) - else: - torch.testing.assert_close( - hf_states.flatten(0, 1), - vllm_states, - rtol=mixtral_moe_tol[dtype], - atol=mixtral_moe_tol[dtype], - ) - - def marlin_moe_generate_valid_test_cases(): import itertools diff --git a/tests/test_regression.py b/tests/test_regression.py index a38b4428dea..385a9286c65 100644 --- a/tests/test_regression.py +++ b/tests/test_regression.py @@ -17,22 +17,6 @@ from vllm import LLM, SamplingParams from vllm.platforms import current_platform -@pytest.mark.skip(reason="In V1, we reject tokens > max_seq_len") -def test_duplicated_ignored_sequence_group(): - """https://github.com/vllm-project/vllm/issues/1655""" - - sampling_params = SamplingParams(temperature=0.01, top_p=0.1, max_tokens=256) - llm = LLM( - model="distilbert/distilgpt2", - max_num_batched_tokens=4096, - tensor_parallel_size=1, - ) - prompts = ["This is a short prompt", "This is a very long prompt " * 1000] - outputs = llm.generate(prompts, sampling_params=sampling_params) - - assert len(prompts) == len(outputs) - - @pytest.mark.parametrize( "model", [ From 55842a8d6961204f5adbcb5ca07da8cf79d85c33 Mon Sep 17 00:00:00 2001 From: Xinyu Chen Date: Sat, 18 Apr 2026 08:53:56 +0800 Subject: [PATCH 099/150] [XPU]fake impl for xpu fp8_gemm (#39984) Signed-off-by: Xinyu Chen Co-authored-by: Kunshang Ji --- vllm/_xpu_ops.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index a0a321173d1..7db074bf920 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -22,6 +22,23 @@ else: except ImportError: from torch.library import impl_abstract as register_fake +if hasattr(torch.ops._xpu_C, "fp8_gemm"): + + @register_fake("_xpu_C::fp8_gemm") + def _fp8_gemm_fake( + q_input: torch.Tensor, + q_weight: torch.Tensor, + out_dtype: torch.dtype, + input_scales: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + input_2d = q_input.view(-1, q_input.shape[-1]) + M = input_2d.size(0) + N = q_weight.size(1) + return torch.empty((M, N), dtype=out_dtype, device=q_input.device) + + if hasattr(torch.ops._xpu_C, "fp8_gemm_w8a16"): @register_fake("_xpu_C::fp8_gemm_w8a16") From 48a65ccb02a19f3b9755c13e8c3a4e8d5a368cc7 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Fri, 17 Apr 2026 22:26:21 -0400 Subject: [PATCH 100/150] [CI] Speed up test_fused_marlin_moe (#40178) Signed-off-by: mgoin Signed-off-by: Michael Goin --- tests/kernels/moe/test_moe.py | 85 +++++++++++++++++++---------------- 1 file changed, 47 insertions(+), 38 deletions(-) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 5cce699dd2c..ebc3256b548 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -143,12 +143,14 @@ MOE_MARLIN_QUANT_TEST_CONFIGS = [ { "a_type": [scalar_types.bfloat16], "b_type": scalar_types.float4_e2m1f, + "c_type": [scalar_types.bfloat16], "group_blocks": [2], }, # MXFP8 { "a_type": [scalar_types.bfloat16], "b_type": scalar_types.float8_e4m3fn, + "c_type": [scalar_types.bfloat16], "group_blocks": [2], }, # AWQ-INT4 with INT8 activation @@ -674,31 +676,35 @@ def test_fused_moe_wn16( torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0) +MARLIN_MOE_SCENARIOS = [ + # (m, n, k, e, topk, ep_size, act_order, is_k_full) + # No act_order: is_k_full=True matches usual case (marlin_is_k_full). + # N>=256 required for Marlin kernel thread config for MXFP8. + # Single token, small matrices + (1, 128, 256, 5, 2, 1, False, True), + # Single token, large matrices + (1, 1024, 2048, 5, 2, 1, False, True), + # Unaligned m, small matrices + (133, 256, 256, 5, 2, 1, False, True), + # Unaligned m, large matrices + (133, 1024, 2048, 12, 3, 1, False, True), + # Aligned batch, small matrices + (128, 256, 256, 5, 2, 1, False, True), + # Aligned batch, large matrices + (128, 1024, 2048, 12, 3, 1, False, True), + # Expert parallelism + (64, 1024, 2048, 12, 3, 4, False, True), + # Act order with is_k_full=True (no tensor parallelism) + (1, 1024, 2048, 5, 2, 1, True, True), + # Act order with is_k_full=False (tensor parallelism) + (133, 256, 256, 5, 2, 1, True, False), +] + + def marlin_moe_generate_valid_test_cases(): import itertools - m_list = [1, 123, 666] - n_list = [128, 1024] - k_list = [256, 2048] - e_list = [5, 12] - topk_list = [2, 3] - ep_size_list = [1, 4] - act_order_list = [True, False] - is_k_full_list = [True, False] - - all_combinations = itertools.product( - MOE_MARLIN_QUANT_TEST_CONFIGS, - m_list, - n_list, - k_list, - e_list, - topk_list, - ep_size_list, - act_order_list, - is_k_full_list, - ) - - def is_invalid( + def is_valid( a_type, b_type, c_type, @@ -715,39 +721,42 @@ def marlin_moe_generate_valid_test_cases(): group_size = group_blocks if group_blocks <= 0 else group_blocks * 16 if group_size > 0 and k % group_size != 0: return False - if act_order and group_size in [-1, k, n]: return False if group_size in [k, n]: return False - if not act_order and is_k_full: + if b_type == scalar_types.float8_e4m3fn and group_size == 32 and is_k_full: return False - return a_type.size_bits < 16 or a_type is c_type cases = [] - for case in all_combinations: - quant_test_config, m, n, k, _, _, _, act_order, *_ = case - if act_order and not quant_test_config.get("support_act_order", False): - continue - + for quant_test_config in MOE_MARLIN_QUANT_TEST_CONFIGS: f16_types = [scalar_types.float16] - inner_combinations = itertools.product( - quant_test_config.get("a_type", f16_types), - [quant_test_config["b_type"]], - quant_test_config.get("c_type", f16_types), - quant_test_config["group_blocks"], + inner_combinations = list( + itertools.product( + quant_test_config.get("a_type", f16_types), + [quant_test_config["b_type"]], + quant_test_config.get("c_type", f16_types), + quant_test_config["group_blocks"], + ) ) + supports_act_order = quant_test_config.get("support_act_order", False) + for sub_case in inner_combinations: if ( sub_case[0] == scalar_types.float8_e4m3fn and current_platform.get_device_capability() not in [89, 120] ): continue - args = sub_case + (m, n, k) + case[4:] - if is_invalid(*args): - cases.append(args) + + for scenario in MARLIN_MOE_SCENARIOS: + m, n, k, e, topk, ep_size, act_order, is_k_full = scenario + if act_order and not supports_act_order: + continue + args = sub_case + (m, n, k, e, topk, ep_size, act_order, is_k_full) + if is_valid(*args): + cases.append(args) return cases From 993859ceb0f4897c423d9fa160fcd80a7a2c890b Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Sat, 18 Apr 2026 10:33:07 +0800 Subject: [PATCH 101/150] [XPU] fix all_reduce all-zero accuracy issue under torch.compile (#39844) Signed-off-by: Chaojun Zhang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/distributed/device_communicators/xpu_communicator.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm/distributed/device_communicators/xpu_communicator.py b/vllm/distributed/device_communicators/xpu_communicator.py index d2e9e89e535..c89c5ddd54a 100644 --- a/vllm/distributed/device_communicators/xpu_communicator.py +++ b/vllm/distributed/device_communicators/xpu_communicator.py @@ -47,9 +47,10 @@ class XpuCommunicator(DeviceCommunicatorBase): self.all2all_manager = AgRsAll2AllManager(self.cpu_group) logger.info("Using AgRs manager on XPU device.") - def all_reduce(self, input_) -> torch.Tensor: - dist.all_reduce(input_, group=self.device_group) - return input_ + def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: + output = input_.clone() if torch.compiler.is_compiling() else input_ + dist.all_reduce(output, group=self.device_group) + return output def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): world_size = self.world_size From b0755523dce2570d74795be2e7462f2ce7d83bfd Mon Sep 17 00:00:00 2001 From: milesial Date: Fri, 17 Apr 2026 20:25:49 -0700 Subject: [PATCH 102/150] [Core] Reduce mm scheduler, get_num_embed overhead (#40143) Signed-off-by: milesial Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- tests/multimodal/test_inputs.py | 9 +++------ vllm/multimodal/inputs.py | 13 ++++++------- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/tests/multimodal/test_inputs.py b/tests/multimodal/test_inputs.py index d6bdf76a6f7..7752a543f42 100644 --- a/tests/multimodal/test_inputs.py +++ b/tests/multimodal/test_inputs.py @@ -26,11 +26,8 @@ def test_placeholder_range_get_num_embeds(is_embed, expected): "is_embed,expected", [ (None, None), - ( - torch.tensor([False, True, False, True, True]), - torch.tensor([0, 1, 1, 2, 3]), - ), - (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])), + (torch.tensor([False, True, False, True, True]), [0, 1, 1, 2, 3]), + (torch.tensor([True, True, True]), [1, 2, 3]), ], ) def test_placeholder_range_embeds_cumsum(is_embed, expected): @@ -41,6 +38,6 @@ def test_placeholder_range_embeds_cumsum(is_embed, expected): assert pr.embeds_cumsum is None return - assert torch.equal(pr.embeds_cumsum, expected) + assert pr.embeds_cumsum == expected # cached_property should return the same object on repeated access assert pr.embeds_cumsum is pr.embeds_cumsum diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index 12356b8727c..d98a1624ac3 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -145,14 +145,15 @@ class PlaceholderRange: """ @cached_property - def embeds_cumsum(self) -> torch.Tensor | None: - return None if self.is_embed is None else self.is_embed.cumsum(dim=0) + def embeds_cumsum(self) -> list[int] | None: + # python list so python indexing avoids torch C++ overhead/conversions/deallocs + return None if self.is_embed is None else self.is_embed.cumsum(dim=0).tolist() def get_num_embeds(self) -> int: if self.embeds_cumsum is None: return self.length - return int(self.embeds_cumsum[-1]) + return self.embeds_cumsum[-1] if self.embeds_cumsum else 0 def get_embeds_indices_in_range( self, start_idx: int, end_idx: int @@ -170,10 +171,8 @@ class PlaceholderRange: if self.embeds_cumsum is None: return start_idx, end_idx - embeds_start_idx = ( - int(self.embeds_cumsum[start_idx - 1]) if start_idx > 0 else 0 - ) - embeds_end_idx = int(self.embeds_cumsum[end_idx - 1]) + embeds_start_idx = self.embeds_cumsum[start_idx - 1] if start_idx > 0 else 0 + embeds_end_idx = self.embeds_cumsum[end_idx - 1] if end_idx > 0 else 0 return embeds_start_idx, embeds_end_idx From d0697cc7b6825da0ba92aff93b05ea85b4725018 Mon Sep 17 00:00:00 2001 From: z1ying <55220715+z1ying@users.noreply.github.com> Date: Fri, 17 Apr 2026 20:26:14 -0700 Subject: [PATCH 103/150] [Doc] Add Realtime Transcription section to supported_models.md (#39845) Signed-off-by: Ziying Tao --- docs/models/supported_models.md | 18 ++++++++++++++++++ docs/serving/openai_compatible_server.md | 2 +- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index f0eb3781fba..746980b8fe1 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -682,6 +682,24 @@ Speech2Text models trained specifically for Automatic Speech Recognition. !!! note `VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed. +#### Realtime Transcription + +Speech models that support streaming transcription via the +[`/v1/realtime`](../serving/openai_compatible_server.md#realtime-api) +WebSocket endpoint. + +| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | +| ------------ | ------ | ----------------- | -------------------- | ------------------------- | +| `VoxtralRealtimeGeneration` | Voxtral Realtime | `mistralai/Voxtral-Mini-4B-Realtime-2602` | | | +| `Qwen3ASRRealtimeGeneration` | Qwen3-ASR Realtime | `Qwen/Qwen3-ASR-0.6B` | | | + +!!! note + `VoxtralRealtimeGeneration` requires `mistral-common[audio]` to be installed, and must be served with `--tokenizer-mode mistral`. + + `Qwen3ASRRealtimeGeneration` is not auto-detected from `config.json`. + You must pass `--hf-overrides '{"architectures":["Qwen3ASRRealtimeGeneration"]}'` + when serving. + ## Pooling Models See [this page](pooling_models/README.md) for more information on how to use pooling models. diff --git a/docs/serving/openai_compatible_server.md b/docs/serving/openai_compatible_server.md index cde7d597d82..a2c90e3abd4 100644 --- a/docs/serving/openai_compatible_server.md +++ b/docs/serving/openai_compatible_server.md @@ -60,7 +60,7 @@ We currently support the following OpenAI APIs: - [Translation API](#translations-api) (`/v1/audio/translations`) - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription). - [Realtime API](#realtime-api) (`/v1/realtime`) - - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription). + - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#realtime-transcription). In addition, we have the following custom APIs: From 80b18230e0ffe94f8198862675a0dc5e86c10997 Mon Sep 17 00:00:00 2001 From: Nithin Chalapathi Date: Fri, 17 Apr 2026 20:31:56 -0700 Subject: [PATCH 104/150] [Frontend] Add multimodal support to /inference/v1/generate endpoint (#38405) Signed-off-by: Nithin Chalapathi Signed-off-by: Nithin Chalapathi Co-authored-by: Cyrus Leung --- .../disaggregated_serving/example_mm_serve.py | 117 +++++++++++++ tests/entrypoints/openai/test_mm_serde.py | 111 ++++++++++++ .../disagg/test_serving_multimodal_tokens.py | 158 ++++++++++++++++++ tests/utils_/test_serial_utils.py | 43 ++++- vllm/entrypoints/serve/disagg/mm_serde.py | 27 +++ vllm/entrypoints/serve/disagg/protocol.py | 17 +- vllm/entrypoints/serve/disagg/serving.py | 48 +++++- vllm/entrypoints/serve/render/serving.py | 24 ++- vllm/utils/serial_utils.py | 22 ++- 9 files changed, 542 insertions(+), 25 deletions(-) create mode 100644 examples/online_serving/disaggregated_serving/example_mm_serve.py create mode 100644 tests/entrypoints/openai/test_mm_serde.py create mode 100644 tests/entrypoints/serve/disagg/test_serving_multimodal_tokens.py create mode 100644 vllm/entrypoints/serve/disagg/mm_serde.py diff --git a/examples/online_serving/disaggregated_serving/example_mm_serve.py b/examples/online_serving/disaggregated_serving/example_mm_serve.py new file mode 100644 index 00000000000..11d81236c57 --- /dev/null +++ b/examples/online_serving/disaggregated_serving/example_mm_serve.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Disaggregated multimodal serving: render → generate round-trip. + +Demonstrates the two-phase disaggregated flow: + 1. /v1/chat/completions/render – preprocesses a multimodal chat request + into token IDs and serialized tensor features. + 2. /inference/v1/generate – runs inference on the preprocessed tokens. + +The render response is passed *directly* to generate with only +``sampling_params`` added, showing that the two endpoints compose with +zero client-side transformation. + +Launch the server first: + + vllm serve Qwen/Qwen3-VL-2B-Instruct \ + --dtype bfloat16 --max-model-len 4096 --enforce-eager + +Then run this script: + + python example_mm_serve.py +""" + +import io + +import pybase64 as base64 +import requests +from PIL import Image +from transformers import AutoTokenizer + +BASE_URL = "http://localhost:8000" +MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct" + + +def make_data_url(image: Image.Image) -> str: + """Encode a PIL image as a base64 data URL.""" + buf = io.BytesIO() + image.save(buf, format="PNG") + b64 = base64.b64encode(buf.getvalue()).decode() + return f"data:image/png;base64,{b64}" + + +def main(): + # -- Step 1: Create a test image (solid red) ------------------------- + image = Image.new("RGB", (224, 224), color=(255, 0, 0)) + data_url = make_data_url(image) + print("Created 224x224 red test image") + + # -- Step 2: Render (preprocess) ------------------------------------- + render_payload = { + "model": MODEL_NAME, + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": data_url}}, + { + "type": "text", + "text": "What color is this image? Answer in one word.", + }, + ], + } + ], + } + + print("\n--- Render ---") + render_resp = requests.post( + f"{BASE_URL}/v1/chat/completions/render", json=render_payload + ) + render_resp.raise_for_status() + render_data = render_resp.json() + + print(f"Response keys: {list(render_data.keys())}") + print(f"Number of token_ids: {len(render_data['token_ids'])}") + + features = render_data.get("features") + if features and features.get("kwargs_data"): + print(f"kwargs_data modalities: {list(features['kwargs_data'].keys())}") + for modality, items in features["kwargs_data"].items(): + print( + f" {modality}: {len(items)} item(s), " + f"first item type: {type(items[0])} length: {len(items[0])}" + if items + else "First item: (empty)" + ) + else: + print("WARNING: no kwargs_data in render response") + + # -- Step 3: Generate (inference) ------------------------------------ + # Pass the render output directly — only add sampling_params. + generate_payload = render_data + generate_payload["sampling_params"] = { + "max_tokens": 20, + "temperature": 0.0, + } + + print("\n--- Generate ---") + gen_resp = requests.post(f"{BASE_URL}/inference/v1/generate", json=generate_payload) + gen_resp.raise_for_status() + gen_data = gen_resp.json() + + # -- Step 4: Decode & print ------------------------------------------ + output_ids = gen_data["choices"][0]["token_ids"] + tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) + text = tokenizer.decode(output_ids, skip_special_tokens=True) + + print(f"Output token count: {len(output_ids)}") + print(f"Generated text: {text!r}") + + if "red" in text.lower(): + print("\nModel correctly identified the red image.") + else: + print(f"\nWARNING: Expected 'red' in output, got: {text!r}") + + +if __name__ == "__main__": + main() diff --git a/tests/entrypoints/openai/test_mm_serde.py b/tests/entrypoints/openai/test_mm_serde.py new file mode 100644 index 00000000000..c568d822e1c --- /dev/null +++ b/tests/entrypoints/openai/test_mm_serde.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Roundtrip tests for multimodal serde used by the disagg generate endpoint.""" + +import torch + +from vllm.entrypoints.serve.disagg.mm_serde import ( + decode_mm_kwargs_item, + encode_mm_kwargs_item, +) +from vllm.entrypoints.serve.disagg.protocol import ( + MultiModalFeatures, + PlaceholderRangeInfo, +) +from vllm.multimodal.inputs import ( + MultiModalBatchedField, + MultiModalFieldElem, + MultiModalFlatField, + MultiModalKwargsItem, + MultiModalSharedField, +) + + +def test_mm_kwargs_item_roundtrip(): + """Full roundtrip test with all three field types and multiple dtypes.""" + e1 = MultiModalFieldElem( + data=torch.zeros(1000, dtype=torch.bfloat16), + field=MultiModalBatchedField(), + ) + e2 = MultiModalFieldElem( + data=torch.ones(100, dtype=torch.int32), + field=MultiModalSharedField(batch_size=4), + ) + e3 = MultiModalFieldElem( + data=torch.randn(20, dtype=torch.float32), + field=MultiModalFlatField(slices=[slice(0, 10), slice(10, 20)], dim=0), + ) + + item = MultiModalKwargsItem({"pixel_values": e1, "grid_thw": e2, "embeds": e3}) + encoded = encode_mm_kwargs_item(item) + + # Encoded result is a base64 string + assert isinstance(encoded, str) + + decoded = decode_mm_kwargs_item(encoded) + + assert set(decoded.keys()) == {"pixel_values", "grid_thw", "embeds"} + assert torch.equal(item["pixel_values"].data, decoded["pixel_values"].data) + assert torch.equal(item["grid_thw"].data, decoded["grid_thw"].data) + assert torch.equal(item["embeds"].data, decoded["embeds"].data) + assert isinstance(decoded["pixel_values"].field, MultiModalBatchedField) + assert isinstance(decoded["grid_thw"].field, MultiModalSharedField) + assert isinstance(decoded["embeds"].field, MultiModalFlatField) + + +def test_mm_kwargs_item_none_data(): + """Roundtrip with None data field.""" + elem = MultiModalFieldElem( + data=None, + field=MultiModalSharedField(batch_size=2), + ) + item = MultiModalKwargsItem({"empty": elem}) + encoded = encode_mm_kwargs_item(item) + decoded = decode_mm_kwargs_item(encoded) + + assert decoded["empty"].data is None + assert isinstance(decoded["empty"].field, MultiModalSharedField) + + +def test_mm_kwargs_item_nested_tensors(): + """Roundtrip with nested tensor data.""" + nested = [torch.randn(3, 4), torch.randn(5, 4)] + elem = MultiModalFieldElem( + data=nested, + field=MultiModalBatchedField(), + ) + item = MultiModalKwargsItem({"nested": elem}) + encoded = encode_mm_kwargs_item(item) + decoded = decode_mm_kwargs_item(encoded) + + decoded_data = decoded["nested"].data + assert len(decoded_data) == 2 + assert torch.equal(nested[0], decoded_data[0]) + assert torch.equal(nested[1], decoded_data[1]) + + +def test_mm_features_with_kwargs_data(): + """Test that MultiModalFeatures can carry serialized tensor data.""" + elem = MultiModalFieldElem( + data=torch.randn(5, 3, dtype=torch.float32), + field=MultiModalBatchedField(), + ) + item = MultiModalKwargsItem({"pixel_values": elem}) + encoded = encode_mm_kwargs_item(item) + + features = MultiModalFeatures( + mm_hashes={"image": ["abc123"]}, + mm_placeholders={"image": [PlaceholderRangeInfo(offset=0, length=10)]}, + kwargs_data={"image": [encoded]}, + ) + + # JSON roundtrip + json_str = features.model_dump_json() + features2 = MultiModalFeatures.model_validate_json(json_str) + + assert features2.mm_hashes == {"image": ["abc123"]} + assert features2.kwargs_data is not None + assert len(features2.kwargs_data["image"]) == 1 + + decoded = decode_mm_kwargs_item(features2.kwargs_data["image"][0]) + assert torch.equal(elem.data, decoded["pixel_values"].data) diff --git a/tests/entrypoints/serve/disagg/test_serving_multimodal_tokens.py b/tests/entrypoints/serve/disagg/test_serving_multimodal_tokens.py new file mode 100644 index 00000000000..e13dd425af1 --- /dev/null +++ b/tests/entrypoints/serve/disagg/test_serving_multimodal_tokens.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for multimodal features through the /inference/v1/generate endpoint. + +Mirrors test_serving_tokens.py but exercises the multimodal piping +using Qwen/Qwen3-VL-2B-Instruct end-to-end via the server's /render -> +/generate -> /detokenize path. Intentionally avoids running the HF +processor in the pytest parent process to keep os.fork() in sibling +tests (e.g. test_weight_transfer_llm.py) deadlock-free. +""" + +import os + +import httpx +import pytest +import pytest_asyncio +from PIL import Image + +from tests.utils import RemoteOpenAIServer +from vllm.multimodal.utils import encode_image_url + +MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct" +GEN_ENDPOINT = "/inference/v1/generate" +RENDER_ENDPOINT = "/v1/chat/completions/render" +DETOKENIZE_ENDPOINT = "/detokenize" + + +@pytest.fixture(scope="module") +def test_image(): + return Image.new("RGB", (224, 224), color=(255, 0, 0)) + + +@pytest.fixture(scope="module") +def server(): + args = [ + "--dtype", + "bfloat16", + "--max-model-len", + "4096", + "--enforce-eager", + "--no-enable-prefix-caching", + ] + + envs = os.environ.copy() + envs["VLLM_ROCM_USE_SKINNY_GEMM"] = "0" + + with RemoteOpenAIServer(MODEL_NAME, args, env_dict=envs) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def client(server: RemoteOpenAIServer): + transport = httpx.AsyncHTTPTransport(uds=server.uds) if server.uds else None + headers = {"Authorization": f"Bearer {server.DUMMY_API_KEY}"} + async with httpx.AsyncClient( + transport=transport, + base_url=server.url_root, + timeout=600, + headers=headers, + ) as c: + yield c + + +@pytest.mark.asyncio +async def test_render_to_generate_roundtrip(client, test_image): + """End-to-end: render a multimodal chat -> feed into generate -> decode. + + All preprocessing and detokenization happens in the server subprocess; + the pytest parent never imports transformers or touches torch tensors. + """ + data_url = encode_image_url(test_image, format="PNG") + + render_payload = { + "model": MODEL_NAME, + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": data_url}}, + { + "type": "text", + "text": "What color is this image? Answer in one word.", + }, + ], + } + ], + } + + render_resp = await client.post(RENDER_ENDPOINT, json=render_payload) + render_resp.raise_for_status() + render_data = render_resp.json() + + # Validate render output structure: keys exist and values are non-empty + # and well-typed. + assert "token_ids" in render_data + assert isinstance(render_data["token_ids"], list) + assert len(render_data["token_ids"]) > 0 + assert all(isinstance(t, int) for t in render_data["token_ids"]) + + assert "features" in render_data + features = render_data["features"] + assert features is not None + assert isinstance(features, dict) + + assert "mm_hashes" in features + assert "image" in features["mm_hashes"] + image_hashes = features["mm_hashes"]["image"] + assert isinstance(image_hashes, list) + assert len(image_hashes) > 0 + assert all(isinstance(h, str) and h for h in image_hashes) + + assert "mm_placeholders" in features + assert "image" in features["mm_placeholders"] + image_placeholders = features["mm_placeholders"]["image"] + assert isinstance(image_placeholders, list) + assert len(image_placeholders) > 0 + for p in image_placeholders: + assert isinstance(p.get("offset"), int) + assert isinstance(p.get("length"), int) + assert p["length"] > 0 + + assert "kwargs_data" in features + assert "image" in features["kwargs_data"] + assert len(features["kwargs_data"]["image"]) > 0 + + # Build generate request from render output + generate_payload = render_data + generate_payload["sampling_params"] = { + "max_tokens": 10, + "temperature": 0.0, + } + + gen_resp = await client.post(GEN_ENDPOINT, json=generate_payload) + gen_resp.raise_for_status() + gen_data = gen_resp.json() + + assert "choices" in gen_data + assert isinstance(gen_data["choices"], list) + assert len(gen_data["choices"]) >= 1 + choice = gen_data["choices"][0] + assert "token_ids" in choice + assert isinstance(choice["token_ids"], list) + assert len(choice["token_ids"]) > 0 + assert all(isinstance(t, int) for t in choice["token_ids"]) + + detok_resp = await client.post( + DETOKENIZE_ENDPOINT, + json={"model": MODEL_NAME, "tokens": choice["token_ids"]}, + ) + detok_resp.raise_for_status() + detok_data = detok_resp.json() + assert "prompt" in detok_data + text = detok_data["prompt"] + assert isinstance(text, str) + assert len(text) > 0 + assert "red" in text.lower(), ( + f"Expected model to identify the red image, got: {text!r}" + ) diff --git a/tests/utils_/test_serial_utils.py b/tests/utils_/test_serial_utils.py index 42e466709cb..85661657d50 100644 --- a/tests/utils_/test_serial_utils.py +++ b/tests/utils_/test_serial_utils.py @@ -7,17 +7,39 @@ from tests.models.utils import check_embeddings_close from vllm.utils.serial_utils import ( EMBED_DTYPES, ENDIANNESS, + MM_METADATA_DTYPES, EmbedDType, Endianness, + MmMetadataDType, binary2tensor, tensor2binary, ) +FLOAT_EMBED_DTYPES = tuple(EMBED_DTYPES.keys()) +INTEGER_EMBED_DTYPES = tuple(MM_METADATA_DTYPES.keys()) + + +def _build_integer_tensor( + embed_dtype: MmMetadataDType, shape: tuple[int, ...] +) -> torch.Tensor: + torch_dtype = MM_METADATA_DTYPES[embed_dtype].torch_dtype + + if torch_dtype is torch.bool: + return torch.randint(0, 2, shape, dtype=torch.int32).to(torch.bool) + if torch_dtype is torch.uint8: + return torch.randint(0, 256, shape, dtype=torch.uint8) + if torch_dtype is torch.int32: + return torch.randint(-(2**20), 2**20, shape, dtype=torch.int32) + if torch_dtype is torch.int64: + return torch.randint(-(2**62), 2**62, shape, dtype=torch.int64) + + raise AssertionError(f"Unsupported non-floating embed dtype: {embed_dtype}") + @pytest.mark.parametrize("endianness", ENDIANNESS) -@pytest.mark.parametrize("embed_dtype", EMBED_DTYPES.keys()) +@pytest.mark.parametrize("embed_dtype", FLOAT_EMBED_DTYPES) @torch.inference_mode() -def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness): +def test_encode_and_decode_floats(embed_dtype: EmbedDType, endianness: Endianness): for i in range(10): tensor = torch.rand(2, 3, 5, 7, 11, 13, device="cpu", dtype=torch.float32) shape = tensor.shape @@ -40,3 +62,20 @@ def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness): name_1="new", tol=1e-2, ) + + +@pytest.mark.parametrize("endianness", ENDIANNESS) +@pytest.mark.parametrize("embed_dtype", INTEGER_EMBED_DTYPES) +@torch.inference_mode() +def test_encode_and_decode_integers( + embed_dtype: MmMetadataDType, endianness: Endianness +): + shape = (2, 3, 5, 7, 11, 13) + + for i in range(10): + tensor = _build_integer_tensor(embed_dtype, shape) + binary = tensor2binary(tensor, embed_dtype, endianness) + new_tensor = binary2tensor(binary, shape, embed_dtype, endianness) + + assert new_tensor.dtype == MM_METADATA_DTYPES[embed_dtype].torch_dtype + torch.testing.assert_close(tensor, new_tensor, atol=0, rtol=0) diff --git a/vllm/entrypoints/serve/disagg/mm_serde.py b/vllm/entrypoints/serve/disagg/mm_serde.py new file mode 100644 index 00000000000..60a3560ad73 --- /dev/null +++ b/vllm/entrypoints/serve/disagg/mm_serde.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Encode/decode utilities for multimodal tensors and field metadata +over JSON/HTTP, used by the disaggregated generate endpoint.""" + +from __future__ import annotations + +import pybase64 + +from vllm.multimodal.inputs import MultiModalKwargsItem +from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder + +_encoder = MsgpackEncoder(size_threshold=2**62) # force all tensors inline +_decoder = MsgpackDecoder(t=MultiModalKwargsItem) + + +def encode_mm_kwargs_item(item: MultiModalKwargsItem) -> str: + """Serialize a MultiModalKwargsItem to a base64 string.""" + bufs = _encoder.encode(item) + assert len(bufs) == 1, "All tensors should be inline" + return pybase64.b64encode(bufs[0]).decode("ascii") + + +def decode_mm_kwargs_item(data: str) -> MultiModalKwargsItem: + """Deserialize a base64 string back to a MultiModalKwargsItem.""" + raw = pybase64.b64decode(data) + return _decoder.decode(raw) diff --git a/vllm/entrypoints/serve/disagg/protocol.py b/vllm/entrypoints/serve/disagg/protocol.py index 345992d3b0b..63369682466 100644 --- a/vllm/entrypoints/serve/disagg/protocol.py +++ b/vllm/entrypoints/serve/disagg/protocol.py @@ -35,14 +35,6 @@ class MultiModalFeatures(BaseModel): Carries hashes (for cache lookup / identification) and placeholder positions so the downstream `/generate` service knows *where* in the token sequence each multimodal item lives. - - Note: - Phase 1 — metadata only. - Phase 2 should add `mm_kwargs` (processed tensor data) using a - binary transport so the ``/generate` side can skip re-processing. - The `/generate` endpoint must also be updated to inject these - features into `EngineInput` before passing to - `InputProcessor.process_inputs`. """ mm_hashes: dict[str, list[str]] @@ -51,6 +43,15 @@ class MultiModalFeatures(BaseModel): mm_placeholders: dict[str, list[PlaceholderRangeInfo]] """Per-modality placeholder ranges in the token sequence.""" + kwargs_data: dict[str, list[str | None]] | None = None + """Per-modality serialized tensor data. + + Each value is a list parallel to ``mm_hashes[modality]``. A ``str`` + entry is a base64-encoded ``MultiModalKwargsItem``; ``None`` means + the item should be resolved from cache. The entire field is + ``None`` for metadata-only (cache-hit) responses. + """ + class GenerateRequest(BaseModel): request_id: str = Field( diff --git a/vllm/entrypoints/serve/disagg/serving.py b/vllm/entrypoints/serve/disagg/serving.py index 14ba85ecf8c..d510125fd03 100644 --- a/vllm/entrypoints/serve/disagg/serving.py +++ b/vllm/entrypoints/serve/disagg/serving.py @@ -25,6 +25,7 @@ from vllm.entrypoints.openai.engine.protocol import ( ) from vllm.entrypoints.openai.engine.serving import OpenAIServing, clamp_prompt_logprobs from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.serve.disagg.mm_serde import decode_mm_kwargs_item from vllm.entrypoints.serve.disagg.protocol import ( GenerateRequest, GenerateResponse, @@ -34,8 +35,14 @@ from vllm.entrypoints.serve.disagg.protocol import ( ) from vllm.entrypoints.serve.render.serving import OpenAIServingRender from vllm.entrypoints.utils import should_include_usage +from vllm.inputs import EngineInput, mm_input from vllm.logger import init_logger from vllm.logprobs import Logprob +from vllm.multimodal.inputs import ( + MultiModalKwargsItem, + MultiModalKwargsItems, + PlaceholderRange, +) from vllm.outputs import RequestOutput from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.utils.collection_utils import as_list @@ -103,11 +110,42 @@ class ServingTokens(OpenAIServing): if raw_request: raw_request.state.request_metadata = request_metadata - (engine_input,) = await self.openai_serving_render.preprocess_completion( - request, - prompt_input=request.token_ids, - prompt_embeds=None, - ) + engine_input: EngineInput + if features := request.features: + # Convert PlaceholderRangeInfo → PlaceholderRange per modality. + mm_placeholders: dict[str, list[PlaceholderRange]] = { + modality: [ + PlaceholderRange(offset=p.offset, length=p.length) for p in ranges + ] + for modality, ranges in features.mm_placeholders.items() + } + + # Deserialize tensor data when present; None → cache hit. + mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {} + if features.kwargs_data is not None: + for modality, items in features.kwargs_data.items(): + mm_kwargs[modality] = [ + decode_mm_kwargs_item(item) if item is not None else None + for item in items + ] + else: + for modality, hashes in features.mm_hashes.items(): + mm_kwargs[modality] = [None] * len(hashes) + + engine_input = mm_input( + prompt_token_ids=request.token_ids, + mm_kwargs=MultiModalKwargsItems(mm_kwargs), + mm_hashes=features.mm_hashes, + mm_placeholders=mm_placeholders, + cache_salt=request.cache_salt, + ) + else: + (engine_input,) = await self.openai_serving_render.preprocess_completion( + request, + prompt_input=request.token_ids, + prompt_embeds=None, + skip_mm_cache=True, + ) # Schedule the request and get the result generator. result_generator: AsyncGenerator[RequestOutput, None] | None = None diff --git a/vllm/entrypoints/serve/render/serving.py b/vllm/entrypoints/serve/render/serving.py index 99b6cb47cb5..3cbf3cc90cc 100644 --- a/vllm/entrypoints/serve/render/serving.py +++ b/vllm/entrypoints/serve/render/serving.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence from http import HTTPStatus -from typing import Any +from typing import Any, cast from openai_harmony import Message as OpenAIMessage @@ -25,6 +25,7 @@ from vllm.entrypoints.openai.parser.harmony_utils import ( render_for_completion, ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.serve.disagg.mm_serde import encode_mm_kwargs_item from vllm.entrypoints.serve.disagg.protocol import ( GenerateRequest, MultiModalFeatures, @@ -37,6 +38,7 @@ from vllm.entrypoints.utils import ( from vllm.inputs import ( EngineInput, MultiModalHashes, + MultiModalInput, MultiModalPlaceholders, PromptType, SingletonPrompt, @@ -251,6 +253,7 @@ class OpenAIServingRender: default_template_kwargs=self.default_chat_template_kwargs, tool_dicts=tool_dicts, tool_parser=tool_parser, + skip_mm_cache=True, reasoning_parser=self.reasoning_parser, ) else: @@ -342,6 +345,7 @@ class OpenAIServingRender: request, prompt_input=request.prompt, prompt_embeds=request.prompt_embeds, + skip_mm_cache=True, ) return engine_inputs @@ -357,9 +361,10 @@ class OpenAIServingRender: if engine_input.get("type") != "multimodal": return None - # At this point engine_input is a MultiModalInputs TypedDict. - mm_hashes: MultiModalHashes = engine_input["mm_hashes"] # type: ignore[typeddict-item] - raw_placeholders: MultiModalPlaceholders = engine_input["mm_placeholders"] # type: ignore[typeddict-item] + # At this point engine_input is a MultiModalInput TypedDict. + mm_engine_input = cast(MultiModalInput, engine_input) + mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"] + raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"] mm_placeholders = { modality: [ @@ -368,9 +373,20 @@ class OpenAIServingRender: for modality, ranges in raw_placeholders.items() } + # Serialize tensor data per modality. + kwargs_data: dict[str, list[str | None]] | None = None + if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"): + kwargs_data = {} + for modality, items in raw_mm_kwargs.items(): + kwargs_data[modality] = [ + encode_mm_kwargs_item(item) if item is not None else None + for item in items + ] + return MultiModalFeatures( mm_hashes=mm_hashes, mm_placeholders=mm_placeholders, + kwargs_data=kwargs_data, ) def _make_request_with_harmony( diff --git a/vllm/utils/serial_utils.py b/vllm/utils/serial_utils.py index 596a7193510..5fde5ac7105 100644 --- a/vllm/utils/serial_utils.py +++ b/vllm/utils/serial_utils.py @@ -27,6 +27,7 @@ class DTypeInfo: EmbedDType = Literal["float32", "float16", "bfloat16", "fp8_e4m3", "fp8_e5m2"] +MmMetadataDType = Literal["int32", "int64", "uint8", "bool"] Endianness = Literal["native", "big", "little"] EncodingFormat = Literal["float", "base64", "bytes", "bytes_only"] @@ -42,6 +43,15 @@ EMBED_DTYPES: Mapping[EmbedDType, DTypeInfo] = { "fp8_e4m3": DTypeInfo(torch.float8_e4m3fn, torch.uint8, np.uint8), "fp8_e5m2": DTypeInfo(torch.float8_e5m2, torch.uint8, np.uint8), } +MM_METADATA_DTYPES: Mapping[MmMetadataDType, DTypeInfo] = { + "int32": DTypeInfo(torch.int32, torch.int32, np.int32), + "int64": DTypeInfo(torch.int64, torch.int64, np.int64), + "uint8": DTypeInfo(torch.uint8, torch.uint8, np.uint8), + "bool": DTypeInfo(torch.bool, torch.uint8, np.uint8), +} +_ALL_SERIAL_DTYPES: Mapping[str, DTypeInfo] = { + k: v for d in (EMBED_DTYPES, MM_METADATA_DTYPES) for k, v in d.items() +} ENDIANNESS: tuple[Endianness, ...] = get_args(Endianness) @@ -56,14 +66,14 @@ def tensor2base64(x: torch.Tensor) -> str: def tensor2binary( tensor: torch.Tensor, - embed_dtype: EmbedDType, + embed_dtype: "EmbedDType | MmMetadataDType", endianness: Endianness, ) -> bytes: assert isinstance(tensor, torch.Tensor) - assert embed_dtype in EMBED_DTYPES + assert embed_dtype in _ALL_SERIAL_DTYPES assert endianness in ENDIANNESS - dtype_info = EMBED_DTYPES[embed_dtype] + dtype_info = _ALL_SERIAL_DTYPES[embed_dtype] np_array = ( tensor.to(dtype_info.torch_dtype) @@ -82,13 +92,13 @@ def tensor2binary( def binary2tensor( binary: bytes, shape: tuple[int, ...], - embed_dtype: EmbedDType, + embed_dtype: "EmbedDType | MmMetadataDType", endianness: Endianness, ) -> torch.Tensor: - assert embed_dtype in EMBED_DTYPES + assert embed_dtype in _ALL_SERIAL_DTYPES assert endianness in ENDIANNESS - dtype_info = EMBED_DTYPES[embed_dtype] + dtype_info = _ALL_SERIAL_DTYPES[embed_dtype] np_array = np.frombuffer(binary, dtype=dtype_info.numpy_view_dtype).reshape(shape) From cda19ecf4db9afba5b03b4013eaa4f52db1656ea Mon Sep 17 00:00:00 2001 From: z1ying <55220715+z1ying@users.noreply.github.com> Date: Fri, 17 Apr 2026 22:31:13 -0700 Subject: [PATCH 105/150] [Doc] Fix outdated source reference comment in anthropic/serving.py (#40189) Signed-off-by: z1ying --- vllm/entrypoints/anthropic/serving.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 5635cfbd359..5136e2b0fb0 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from -# https://github.com/vllm/vllm/entrypoints/openai/serving_chat.py +# https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/chat_completion/serving.py """Anthropic Messages API serving handler""" From aeee7ef9391028939afd08e20c12a1e279efbdf1 Mon Sep 17 00:00:00 2001 From: Rishapveer Singh Date: Sat, 18 Apr 2026 07:34:33 +0200 Subject: [PATCH 106/150] [Bugfix] Fix k_proj's bias for GLM-ASR (#40160) Signed-off-by: Rishapveer Singh --- vllm/model_executor/models/glmasr.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/glmasr.py b/vllm/model_executor/models/glmasr.py index 0d54588cc2b..96be79a3093 100644 --- a/vllm/model_executor/models/glmasr.py +++ b/vllm/model_executor/models/glmasr.py @@ -66,7 +66,7 @@ from .interfaces import ( SupportsTranscription, ) from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix -from .whisper import ISO639_1_SUPPORTED_LANGS +from .whisper import ISO639_1_SUPPORTED_LANGS, _create_fake_bias_for_k_proj class GlmAsrEncoderRotaryEmbedding(nn.Module): @@ -499,6 +499,8 @@ class GlmAsrEncoder(nn.Module): """Custom weight loading to handle q_proj/k_proj/v_proj -> qkv_proj mapping.""" from vllm.model_executor.model_loader.weight_utils import default_weight_loader + weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight") + stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), From 87518c302797be0b20bd720fb3a7e195a517bc66 Mon Sep 17 00:00:00 2001 From: Chinmay-Kulkarni-AMD Date: Sat, 18 Apr 2026 11:52:37 +0530 Subject: [PATCH 107/150] [ZenCPU] AMD Zen CPU Backend with supported dtypes via zentorch weekly (#39967) Signed-off-by: Chinmay Kulkarni --- setup.py | 4 +++- vllm/platforms/zen_cpu.py | 8 ++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 7665978c6b0..f6276616a19 100644 --- a/setup.py +++ b/setup.py @@ -1085,7 +1085,9 @@ setup( install_requires=get_requirements(), extras_require={ # AMD Zen CPU optimizations via zentorch - "zen": ["zentorch"], + "zen": [ + "zentorch-weekly==5.2.1.dev20260408" + ], # Zentorch has weekly releases. This pulls the known-good version. "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"], "tensorizer": ["tensorizer==2.10.1"], "fastsafetensors": ["fastsafetensors >= 0.2.2"], diff --git a/vllm/platforms/zen_cpu.py b/vllm/platforms/zen_cpu.py index 481eec1cb4e..2af64e5e9f5 100644 --- a/vllm/platforms/zen_cpu.py +++ b/vllm/platforms/zen_cpu.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + from vllm.logger import init_logger from vllm.platforms.cpu import CpuPlatform @@ -22,3 +24,9 @@ class ZenCpuPlatform(CpuPlatform): def is_zen_cpu(self) -> bool: # is_cpu() also returns True for this platform (inherited from CpuPlatform). return True + + # Currently, AMD CPUs do not support float16 compute. + # Hence explicitly return bfloat16 and float32. + @property + def supported_dtypes(self) -> list[torch.dtype]: + return [torch.bfloat16, torch.float32] From 153ba7f0f3d882c4feb84cea63b6a4e9618c9537 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Sat, 18 Apr 2026 02:55:38 -0400 Subject: [PATCH 108/150] [Refactor] Drop direct dependency on librosa (#39079) Signed-off-by: Nick Cao Co-authored-by: Claude --- docs/contributing/model/transcription.md | 6 ++--- docs/features/multimodal_inputs.md | 6 ++--- ...i_chat_completion_client_for_multimodal.py | 6 ++--- .../online_serving/openai_realtime_client.py | 5 ++--- .../test_transcription_api_correctness.py | 6 ++--- .../realtime/test_realtime_validation.py | 4 ++-- .../test_transcription_validation_whisper.py | 4 ++-- .../test_translation_validation.py | 4 ++-- .../multimodal/generation/test_phi4mm.py | 4 ++-- .../multimodal/generation/test_whisper.py | 7 +++--- tests/multimodal/media/test_audio.py | 4 ++-- vllm/multimodal/media/audio.py | 6 ++--- .../processors/cohere_asr.py | 22 +++++++++---------- 13 files changed, 40 insertions(+), 44 deletions(-) diff --git a/docs/contributing/model/transcription.md b/docs/contributing/model/transcription.md index a23de100da3..db868686e2a 100644 --- a/docs/contributing/model/transcription.md +++ b/docs/contributing/model/transcription.md @@ -193,7 +193,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics: The API server takes care of basic audio I/O and optional chunking before building prompts: -- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `librosa`. +- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`. - Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`. - Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words. @@ -206,8 +206,8 @@ Relevant server logic: async def _preprocess_speech_to_text(...): language = self.model_cls.validate_language(request.language) ... - y, sr = librosa.load(bytes_, sr=self.asr_config.sample_rate) - duration = librosa.get_duration(y=y, sr=sr) + y, sr = load_audio(bytes_, sr=self.asr_config.sample_rate) + duration = get_audio_duration(y=y, sr=sr) do_split_audio = (self.asr_config.allow_audio_chunking and duration > self.asr_config.max_audio_clip_s) chunks = [y] if not do_split_audio else self._split_audio(y, int(sr)) diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index ee82c34fa0e..d9b49f7cb7f 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -300,12 +300,12 @@ Full example: [examples/offline_inference/audio_language.py](../../examples/offl Speech-to-text models like Whisper have a maximum audio length they can process (typically 30 seconds). For longer audio files, vLLM provides a utility to intelligently split audio into chunks at quiet points to minimize cutting through speech. ```python -import librosa from vllm import LLM, SamplingParams from vllm.multimodal.audio import split_audio +from vllm.multimodal.media.audio import load_audio # Load long audio file -audio, sr = librosa.load("long_audio.wav", sr=16000) +audio, sr = load_audio("long_audio.wav", sr=16000) # Split into chunks at low-energy (quiet) regions chunks = split_audio( @@ -832,7 +832,7 @@ Then, you can use the OpenAI client as follows: base_url=openai_api_base, ) - # Any format supported by librosa is supported + # Any format supported by soundfile/PyAV is supported audio_url = AudioAsset("winning_call").url audio_base64 = encode_base64_content_from_url(audio_url) diff --git a/examples/online_serving/openai_chat_completion_client_for_multimodal.py b/examples/online_serving/openai_chat_completion_client_for_multimodal.py index c4407923ed2..ba3adf55c90 100644 --- a/examples/online_serving/openai_chat_completion_client_for_multimodal.py +++ b/examples/online_serving/openai_chat_completion_client_for_multimodal.py @@ -267,7 +267,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None: { "type": "input_audio", "input_audio": { - # Any format supported by librosa is supported + # Any format supported by soundfile/PyAV is supported "data": audio_base64, "format": "wav", }, @@ -292,7 +292,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None: { "type": "audio_url", "audio_url": { - # Any format supported by librosa is supported + # Any format supported by soundfile/PyAV is supported "url": audio_url }, }, @@ -316,7 +316,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None: { "type": "audio_url", "audio_url": { - # Any format supported by librosa is supported + # Any format supported by soundfile/PyAV is supported "url": f"data:audio/ogg;base64,{audio_base64}" }, }, diff --git a/examples/online_serving/openai_realtime_client.py b/examples/online_serving/openai_realtime_client.py index 2bd3c7e60d5..fda3d7cb456 100644 --- a/examples/online_serving/openai_realtime_client.py +++ b/examples/online_serving/openai_realtime_client.py @@ -12,7 +12,6 @@ model, for example: Requirements: - vllm with audio support - websockets -- librosa - numpy The script: @@ -26,12 +25,12 @@ import argparse import asyncio import json -import librosa import numpy as np import pybase64 as base64 import websockets from vllm.assets.audio import AudioAsset +from vllm.multimodal.media.audio import load_audio def audio_to_pcm16_base64(audio_path: str) -> str: @@ -39,7 +38,7 @@ def audio_to_pcm16_base64(audio_path: str) -> str: Load an audio file and convert it to base64-encoded PCM16 @ 16kHz. """ # Load audio and resample to 16kHz mono - audio, _ = librosa.load(audio_path, sr=16000, mono=True) + audio, _ = load_audio(audio_path, sr=16000, mono=True) # Convert to PCM16 pcm16 = (audio * 32767).astype(np.int16) # Encode as base64 diff --git a/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py b/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py index 194c52eae35..f847b7b0d88 100644 --- a/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py +++ b/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py @@ -13,7 +13,6 @@ import io import time from statistics import mean, median -import librosa import pytest import soundfile import torch @@ -21,6 +20,7 @@ from datasets import load_dataset from evaluate import load from transformers.models.whisper.english_normalizer import EnglishTextNormalizer +from vllm.multimodal.audio import get_audio_duration from vllm.tokenizers import get_tokenizer from ....models.registry import HF_EXAMPLE_MODELS @@ -84,7 +84,7 @@ async def process_dataset(model, client, data, concurrent_request): trust_remote_code=model_info.trust_remote_code, ) - # Warmup call as the first `librosa.load` server-side is quite slow. + # Warmup call as the first `load_audio` server-side is quite slow. audio, sr = data[0]["audio"]["array"], data[0]["audio"]["sampling_rate"] _ = await bound_transcribe(sem, client, tokenizer, (audio, sr), "") @@ -118,7 +118,7 @@ def print_performance_metrics(results, total_time): def add_duration(sample): y, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"] - sample["duration_ms"] = librosa.get_duration(y=y, sr=sr) * 1000 + sample["duration_ms"] = get_audio_duration(y=y, sr=sr) * 1000 return sample diff --git a/tests/entrypoints/openai/realtime/test_realtime_validation.py b/tests/entrypoints/openai/realtime/test_realtime_validation.py index bb6b02f5c99..e317090fa54 100644 --- a/tests/entrypoints/openai/realtime/test_realtime_validation.py +++ b/tests/entrypoints/openai/realtime/test_realtime_validation.py @@ -5,7 +5,6 @@ import asyncio import json import warnings -import librosa import numpy as np import pybase64 as base64 import pytest @@ -14,6 +13,7 @@ import websockets from tests.entrypoints.openai.conftest import add_attention_backend from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer from vllm.assets.audio import AudioAsset +from vllm.multimodal.media.audio import load_audio # Increase engine iteration timeout for ROCm where first-use JIT compilation # can exceed the default 60s, causing a silent deadlock in feed_tokens. @@ -56,7 +56,7 @@ async def send_event(ws, event: dict) -> None: def mary_had_lamb_audio_chunks() -> list[str]: """Audio split into ~1 second chunks for streaming.""" path = AudioAsset("mary_had_lamb").get_local_path() - audio, _ = librosa.load(str(path), sr=16000, mono=True) + audio, _ = load_audio(str(path), sr=16000, mono=True) # Split into ~0.1 second chunks (1600 samples at 16kHz) chunk_size = 1600 diff --git a/tests/entrypoints/openai/speech_to_text/test_transcription_validation_whisper.py b/tests/entrypoints/openai/speech_to_text/test_transcription_validation_whisper.py index 8dba1b59742..511179f7fcb 100644 --- a/tests/entrypoints/openai/speech_to_text/test_transcription_validation_whisper.py +++ b/tests/entrypoints/openai/speech_to_text/test_transcription_validation_whisper.py @@ -6,7 +6,6 @@ import asyncio import io import json -import librosa import numpy as np import openai import pytest @@ -14,6 +13,7 @@ import pytest_asyncio import soundfile as sf from tests.utils import RemoteOpenAIServer +from vllm.multimodal.media.audio import load_audio from vllm.platforms import current_platform MODEL_NAME = "openai/whisper-large-v3-turbo" @@ -134,7 +134,7 @@ async def test_bad_requests(mary_had_lamb, whisper_client): @pytest.mark.asyncio async def test_long_audio_request(mary_had_lamb, whisper_client): mary_had_lamb.seek(0) - audio, sr = librosa.load(mary_had_lamb) + audio, sr = load_audio(mary_had_lamb) # Add small silence after each audio for repeatability in the split process audio = np.pad(audio, (0, 1600)) repeated_audio = np.tile(audio, 10) diff --git a/tests/entrypoints/openai/speech_to_text/test_translation_validation.py b/tests/entrypoints/openai/speech_to_text/test_translation_validation.py index 16b9614d957..a8b17bf3432 100644 --- a/tests/entrypoints/openai/speech_to_text/test_translation_validation.py +++ b/tests/entrypoints/openai/speech_to_text/test_translation_validation.py @@ -7,7 +7,6 @@ import io import json import httpx -import librosa import numpy as np import openai import pytest @@ -17,6 +16,7 @@ import soundfile as sf from tests.entrypoints.openai.conftest import add_attention_backend from tests.utils import RemoteOpenAIServer from vllm.logger import init_logger +from vllm.multimodal.media.audio import load_audio logger = init_logger(__name__) @@ -264,7 +264,7 @@ async def test_long_audio_request(foscolo, client_and_model): if model_name == "google/gemma-3n-E2B-it": pytest.skip("Gemma3n does not support long audio requests") foscolo.seek(0) - audio, sr = librosa.load(foscolo) + audio, sr = load_audio(foscolo) repeated_audio = np.tile(audio, 2) # Repeated audio to buffer buffer = io.BytesIO() diff --git a/tests/models/multimodal/generation/test_phi4mm.py b/tests/models/multimodal/generation/test_phi4mm.py index 7f1a12f0447..1a4fb35a28a 100644 --- a/tests/models/multimodal/generation/test_phi4mm.py +++ b/tests/models/multimodal/generation/test_phi4mm.py @@ -4,7 +4,6 @@ import os from collections.abc import Sequence -import librosa import pytest import regex as re from huggingface_hub import snapshot_download @@ -14,6 +13,7 @@ from vllm.assets.image import ImageAsset from vllm.logprobs import SampleLogprobs from vllm.lora.request import LoRARequest from vllm.multimodal.image import convert_image_mode, rescale_image_size +from vllm.multimodal.media.audio import load_audio from ....conftest import ( IMAGE_ASSETS, @@ -290,7 +290,7 @@ def test_vision_speech_models( num_logprobs: int, ) -> None: # use the example speech question so that the model outputs are reasonable - audio = librosa.load(speech_question, sr=None) + audio = load_audio(speech_question, sr=None) image = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB") inputs_vision_speech = [ diff --git a/tests/models/multimodal/generation/test_whisper.py b/tests/models/multimodal/generation/test_whisper.py index babf7e7a497..186e7e054ce 100644 --- a/tests/models/multimodal/generation/test_whisper.py +++ b/tests/models/multimodal/generation/test_whisper.py @@ -4,11 +4,11 @@ from collections.abc import Sequence from typing import Any -import librosa import pytest from transformers import AutoModelForSpeechSeq2Seq from vllm.assets.audio import AudioAsset +from vllm.multimodal.audio import AudioResampler from vllm.platforms import current_platform from ....conftest import HfRunner, PromptAudioInput, VllmRunner @@ -93,13 +93,12 @@ def run_test( def resampled_assets() -> list[tuple[Any, int]]: audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")] sampled_assets = [] + resampler = AudioResampler(target_sr=WHISPER_SAMPLE_RATE) for asset in audio_assets: audio, orig_sr = asset.audio_and_sample_rate # Resample to Whisper's expected sample rate (16kHz) if orig_sr != WHISPER_SAMPLE_RATE: - audio = librosa.resample( - audio, orig_sr=orig_sr, target_sr=WHISPER_SAMPLE_RATE - ) + audio = resampler.resample(audio, orig_sr=orig_sr) sampled_assets.append( (audio, WHISPER_SAMPLE_RATE), ) diff --git a/tests/multimodal/media/test_audio.py b/tests/multimodal/media/test_audio.py index 4361066ab88..3729e71f24e 100644 --- a/tests/multimodal/media/test_audio.py +++ b/tests/multimodal/media/test_audio.py @@ -3,12 +3,12 @@ from pathlib import Path from unittest.mock import patch -import librosa import numpy as np import pybase64 as base64 import pytest from vllm.multimodal.media import AudioMediaIO +from vllm.multimodal.media.audio import load_audio from ...conftest import AudioTestAssets @@ -73,6 +73,6 @@ def test_audio_media_io_from_video(video_assets): video_path = video_assets[0].video_path with open(video_path, "rb") as f: audio, sr = audio_io.load_bytes(f.read()) - audio_ref, sr_ref = librosa.load(video_path, sr=None) + audio_ref, sr_ref = load_audio(video_path, sr=None) assert sr == sr_ref np.testing.assert_allclose(audio_ref, audio, atol=1e-4) diff --git a/vllm/multimodal/media/audio.py b/vllm/multimodal/media/audio.py index 26d58d13a73..37b8662a76b 100644 --- a/vllm/multimodal/media/audio.py +++ b/vllm/multimodal/media/audio.py @@ -29,9 +29,9 @@ except ImportError: soundfile = PlaceholderModule("soundfile") # type: ignore[assignment] -# Public libsndfile error codes exposed via `soundfile.LibsndfileError.code`, soundfile -# being librosa's main backend. Used to validate if an audio loading error is due to a -# server error vs a client error (invalid audio file). +# Public libsndfile error codes exposed via `soundfile.LibsndfileError.code`, +# soundfile being the main audio loading backend. Used to validate if an audio +# loading error is due to a server error vs a client error (invalid audio file). # 0 = sf_error(NULL) race condition: when multiple threads fail sf_open_virtual # concurrently, one thread may clear the global error before another reads it, # producing code=0 ("Garbled error message from libsndfile" in soundfile). diff --git a/vllm/transformers_utils/processors/cohere_asr.py b/vllm/transformers_utils/processors/cohere_asr.py index f742074a4e3..e1257de4e73 100644 --- a/vllm/transformers_utils/processors/cohere_asr.py +++ b/vllm/transformers_utils/processors/cohere_asr.py @@ -4,11 +4,11 @@ import logging import math import random -import librosa import numpy as np import torch import torch.nn.functional as F from torch import nn +from torchaudio.functional import melscale_fbanks from transformers import AutoFeatureExtractor, AutoProcessor, BatchFeature from transformers.feature_extraction_sequence_utils import ( SequenceFeatureExtractor, @@ -129,17 +129,15 @@ class FilterbankFeatures(nn.Module): self.pad_min_duration = 0.0 self.pad_direction = "both" - filterbanks = torch.tensor( - librosa.filters.mel( - sr=sample_rate, - n_fft=self.n_fft, - n_mels=nfilt, - fmin=lowfreq, - fmax=highfreq, - norm=mel_norm, - ), - dtype=torch.float, - ).unsqueeze(0) + filterbanks = melscale_fbanks( + n_freqs=self.n_fft // 2 + 1, + f_min=lowfreq, + f_max=highfreq, + n_mels=nfilt, + sample_rate=sample_rate, + norm=mel_norm, + mel_scale="slaney", + ).T.unsqueeze(0) self.register_buffer("fb", filterbanks) # Calculate maximum sequence length From bfde49e287cb5522fb0625c8e2b4e03cac20cbb2 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Sat, 18 Apr 2026 15:41:37 +0800 Subject: [PATCH 109/150] [DOC] Add fuse_minimax_qk_norm (#39782) Signed-off-by: Jee Jee Li --- docs/design/fusions.md | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 3fe1f769bb7..1df0eb6f439 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -21,6 +21,7 @@ or just on the low or high end. | Fusion | `PassConfig` flag | Fused operations | Default at | E2E Speedup | Fullgraph | `num_tokens` | | ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ | | [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms) | `fuse_allreduce_rms` | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20% | No | Low | +| [MiniMax QK Norm](#minimax-qk-norm-fuse_minimax_qk_norm) | `fuse_minimax_qk_norm` | Q/K variance all-reduce → Q/K RMSNorm | Off by default | 2-3% | No | Low | | [Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | Attention output → FP8/NVFP4 quant | Off by default | 3-7% | Yes | Always | | [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always | | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low | @@ -40,6 +41,7 @@ The table below lists the quantization schemes supported by each fusion on each | Fusion | SM100 (Blackwell) | SM90 (Hopper) | SM89 (Ada) | SM80 (Ampere) | ROCm | | ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- | | `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — | +| `fuse_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — | | `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* | | `fuse_attn_quant` (MLA)\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static(untested)\* | | `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 | @@ -54,6 +56,9 @@ The table below lists the quantization schemes supported by each fusion on each fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant) for per-backend details. +\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires +tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`. + † `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today; other architectures support requires setting `PassConfig.sp_min_token_num` explicitly. SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`. @@ -184,6 +189,35 @@ If these conditions are set, the fusion is enabled automatically for optimizatio - Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py) +### MiniMax QK Norm (`fuse_minimax_qk_norm`) + +!!! info + This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold: + the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`), + and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any + optimization level. + +**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the +per-token Q/K variances before applying RMS normalization to Q and K. + +This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion): +`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while +`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models. + +Example: + +```bash +vllm serve MiniMaxAI/MiniMax-M2.5 \ + --tensor-parallel-size 4 \ + --compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}' +``` + +**Code locations.** + +- Pass: [`vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py) +- CUDA op: [`csrc/minimax_reduce_rms_kernel.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/minimax_reduce_rms_kernel.cu) (`minimax_allreduce_rms_qk`) +- Workspace helper: [`vllm/model_executor/layers/mamba/lamport_workspace.py`](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/lamport_workspace.py) + ### Sequence Parallelism (`enable_sp`) **What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather, From b5f6c5f8343dd0ff5accb3f8a8f7d1f5c40f32aa Mon Sep 17 00:00:00 2001 From: Yusuf Mohammad <79484377+YM2132@users.noreply.github.com> Date: Sat, 18 Apr 2026 15:05:21 +0100 Subject: [PATCH 110/150] Added general ND x ND matmul and unit test for it (#39909) Signed-off-by: Yusuf --- .../test_matmul_batch_invariant.py | 105 ++++++++++++++++++ vllm/model_executor/layers/batch_invariant.py | 55 ++++----- 2 files changed, 129 insertions(+), 31 deletions(-) create mode 100644 tests/v1/determinism/test_matmul_batch_invariant.py diff --git a/tests/v1/determinism/test_matmul_batch_invariant.py b/tests/v1/determinism/test_matmul_batch_invariant.py new file mode 100644 index 00000000000..8838b6be270 --- /dev/null +++ b/tests/v1/determinism/test_matmul_batch_invariant.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Test batch-invariant matmul against torch.matmul for various shape combinations. + +Tests correctness (matches torch.matmul) and batch invariance (result for one +item doesn't change based on other items in the batch). +""" + +import pytest +import torch +from utils import skip_unsupported + +from vllm.model_executor.layers.batch_invariant import matmul_batch_invariant +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type + + +@skip_unsupported +@pytest.mark.parametrize( + "a_shape,b_shape", + [ + # 2D x 2D + ((32, 64), (64, 16)), + # 2D x 3D + ((64, 16), (4, 16, 32)), + # 3D x 2D + ((4, 32, 64), (64, 16)), + # 4D x 2D + ((1, 4, 32, 64), (64, 16)), + # 3D x 3D + ((4, 32, 64), (4, 64, 16)), + # 3D x 4D + ((2, 32, 64), (1, 2, 64, 16)), + # 4D x 3D (Gemma4 pattern) + ((1, 2, 32, 64), (2, 64, 16)), + # 4D x 4D + ((1, 2, 32, 64), (4, 2, 64, 16)), + # 2D x 4D + ((32, 64), (1, 2, 64, 16)), + # 2D x 5D + ((32, 64), (1, 2, 2, 64, 16)), + # 5D x 2D + ((1, 2, 2, 32, 64), (64, 16)), + # 5D x 5D + ((1, 2, 4, 32, 64), (1, 2, 4, 64, 16)), + ], +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_matmul_correctness(a_shape, b_shape, dtype): + """ + Compare matmul_batch_invariant against torch.matmul for various shapes. + """ + device = torch.device(DEVICE_TYPE) + + torch.manual_seed(42) + a = torch.rand(a_shape, dtype=dtype, device=device) + b = torch.rand(b_shape, dtype=dtype, device=device) + + # Standard implementation (CUDA ops) + standard_output = torch.matmul(a, b) + + # Batch-invariant implementation (Triton) + triton_output = matmul_batch_invariant(a, b) + + # Compare outputs + # Use looser tolerance for bfloat16 due to its lower precision + if dtype == torch.bfloat16: + rtol, atol = 1e-1, 1e-1 # 10% relative tolerance for bfloat16 + else: + rtol, atol = 1e-2, 1e-2 # 1% for float16/float32 + + torch.testing.assert_close( + triton_output, + standard_output, + rtol=rtol, + atol=atol, + msg=f"matmul mismatch for a ndim={a.ndim}, b ndim={b.ndim},", + ) + + +@skip_unsupported +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_matmul_batch_invariance(dtype): + """ + Verify that the result for one item is bitwise identical regardless + of what other items are in the batch. + """ + + device = torch.device(DEVICE_TYPE) + + torch.manual_seed(42) + a_single = torch.rand((1, 64, 32), dtype=dtype, device=device) + b = torch.rand((32, 128), dtype=dtype, device=device) + + standard_output = matmul_batch_invariant(a_single, b) + + a_batch = torch.rand((8, 64, 32), dtype=dtype, device=device) + a_batch[3] = a_single[0] + + batch_output = matmul_batch_invariant(a_batch, b) + batch_output_a = batch_output[3] + + assert torch.equal(standard_output[0], batch_output_a) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index 98fd6be8f6a..4a88421e3b5 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math import os from collections.abc import Callable from typing import Any @@ -611,51 +612,43 @@ def matmul_batch_invariant(a, b, *, out=None): out.copy_(result) return out return result - elif a.ndim == 3 and b.ndim == 3: - # Handle batched case like bmm - return bmm_batch_invariant(a, b, out=out) - elif a.ndim == 3 and b.ndim == 2: - # Handle 3D x 2D: common for linear layers - # (batch, seq, hidden) @ (hidden, out) -> (batch, seq, out) - # Reshape to 2D, do mm, reshape back - batch, seq, hidden = a.shape + elif b.ndim == 2: + # Handle ND x 2D: Common for linear layers + # (..., batch, seq, hidden) @ (hidden, out) -> (..., batch, seq, out) + batch_dims = a.shape[:-1] + hidden = a.shape[-1] + out_dim = b.shape[-1] a_2d = a.reshape(-1, hidden) result_2d = matmul_persistent(a_2d, b) - result = result_2d.reshape(batch, seq, -1) + result = result_2d.reshape(batch_dims + (out_dim,)) if out is not None: out.copy_(result) return out return result - elif a.ndim == 2 and b.ndim == 3: - # Handle 2D x 3D: (M, K) @ (B, K, N) -> (B, M, N) - # By broadcasting `a` to 3D, we can reuse the batched matrix - # multiplication logic. - a_expanded = a.unsqueeze(0).expand(b.shape[0], -1, -1) - return bmm_batch_invariant(a_expanded, b, out=out) - elif a.ndim == 4 and b.ndim == 4: - # Handle 4D attention tensors: [batch, heads, seq, dim] - # Reshape to 3D, process, reshape back - batch, heads, seq_a, dim_a = a.shape - _, _, dim_b, seq_b = b.shape - - # Reshape to [batch*heads, seq_a, dim_a] - a_3d = a.reshape(batch * heads, seq_a, dim_a) - b_3d = b.reshape(batch * heads, dim_b, seq_b) - + elif a.ndim >= 2 and b.ndim >= 3: + # Generic handler for 2D x ND and ND x ND (except 1D) + # Broadcast dims to ensure both matrices have the same shape + # If 2D x ND, then unsqueeze to add a dim to a + if a.ndim == 2: + a = a.unsqueeze(0) + broadcast_shape = torch.broadcast_shapes(a.shape[:-2], b.shape[:-2]) + a = a.expand(broadcast_shape + a.shape[-2:]) + b = b.expand(broadcast_shape + b.shape[-2:]) + batch_dim = math.prod(broadcast_shape) + # Reuse broadcast shape to get all dims except mm dims + a_3d = a.reshape(batch_dim, a.shape[-2], a.shape[-1]) + b_3d = b.reshape(batch_dim, b.shape[-2], b.shape[-1]) # Do batched matmul result_3d = bmm_batch_invariant(a_3d, b_3d) - - # Reshape back to [batch, heads, seq_a, seq_b] - result = result_3d.reshape(batch, heads, seq_a, seq_b) - + # Reshape back to [broadcast_shape, seq_a, seq_b] + result = result_3d.reshape(broadcast_shape + (a.shape[-2], b.shape[-1])) if out is not None: out.copy_(result) return out return result else: raise ValueError( - f"matmul_batch_invariant currently only supports 2D x 2D, 3D x 3D, " - f"3D x 2D, 2D x 3D, and 4D x 4D, " + f"matmul_batch_invariant requires both inputs be at least 2D " f"got shapes {a.shape} and {b.shape}" ) From ed0622e3a809fe399b81009c9dbb9cf7299414e6 Mon Sep 17 00:00:00 2001 From: Dan Alistarh Date: Sat, 18 Apr 2026 20:31:59 +0200 Subject: [PATCH 111/150] [Attention] TurboQuant: remove redundant random signs, add prior art attribution (#40194) Signed-off-by: Dan Alistarh --- tests/quantization/test_turboquant.py | 56 +++++-------------- .../layers/attention/attention.py | 19 +------ .../quantization/turboquant/__init__.py | 17 ++++-- .../layers/quantization/turboquant/config.py | 24 +++++--- .../quantization/turboquant/quantizer.py | 18 ------ vllm/v1/attention/backends/turboquant_attn.py | 22 +++----- 6 files changed, 53 insertions(+), 103 deletions(-) diff --git a/tests/quantization/test_turboquant.py b/tests/quantization/test_turboquant.py index 519022346c2..90beb64a474 100644 --- a/tests/quantization/test_turboquant.py +++ b/tests/quantization/test_turboquant.py @@ -18,9 +18,6 @@ from vllm.model_executor.layers.quantization.turboquant.config import ( TQ_PRESETS, TurboQuantConfig, ) -from vllm.model_executor.layers.quantization.turboquant.quantizer import ( - generate_wht_signs, -) from vllm.platforms import current_platform from vllm.utils.math_utils import next_power_of_2 @@ -393,7 +390,7 @@ class TestRotationMatrix: # ============================================================================ -# WHT rotation tests (serving path: generate_wht_signs + _build_hadamard) +# Hadamard rotation tests (serving path: _build_hadamard) # ============================================================================ @@ -406,50 +403,26 @@ def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor: @pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available") -class TestWHTRotation: - """Tests for the WHT rotation actually used in serving.""" +class TestHadamardRotation: + """Tests for the Hadamard rotation used in serving.""" @pytest.mark.parametrize("dim", [64, 128, 256]) - def test_wht_orthonormal(self, dim): - """signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I.""" - signs = generate_wht_signs(dim, seed=42, device=DEVICE_TYPE) + def test_hadamard_orthonormal(self, dim): + """H must be orthonormal: H @ H^T = I.""" H = _build_hadamard(dim, DEVICE_TYPE) - PiT = (signs.unsqueeze(1) * H).contiguous() - eye = PiT @ PiT.T + eye = H @ H.T assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), ( - f"WHT rotation not orthonormal for dim={dim}" + f"Hadamard not orthonormal for dim={dim}" ) @pytest.mark.parametrize("dim", [64, 128, 256]) - def test_wht_self_inverse(self, dim): - """PiT should be self-inverse: PiT @ PiT = I (up to sign flip).""" - signs = generate_wht_signs(dim, seed=42, device=DEVICE_TYPE) + def test_hadamard_symmetric(self, dim): + """Sylvester Hadamard must be symmetric: H = H^T.""" H = _build_hadamard(dim, DEVICE_TYPE) - PiT = (signs.unsqueeze(1) * H).contiguous() - Pi = PiT.T.contiguous() - # Pi @ PiT should be identity (rotation then inverse) - result = Pi @ PiT - assert torch.allclose(result, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), ( - f"WHT rotation not self-inverse for dim={dim}" + assert torch.allclose(H, H.T, atol=1e-6), ( + f"Hadamard not symmetric for dim={dim}" ) - def test_wht_signs_deterministic(self): - """Same seed must produce identical signs.""" - s1 = generate_wht_signs(128, seed=42) - s2 = generate_wht_signs(128, seed=42) - assert torch.equal(s1, s2) - - def test_wht_signs_different_seeds(self): - """Different seeds must produce different signs.""" - s1 = generate_wht_signs(128, seed=42) - s2 = generate_wht_signs(128, seed=99) - assert not torch.equal(s1, s2) - - def test_wht_signs_are_pm1(self): - """All sign values must be exactly +1 or -1.""" - signs = generate_wht_signs(128, seed=42) - assert torch.all(signs.abs() == 1.0) - # ============================================================================ # Store → Decode round-trip test (GPU + Triton required) @@ -491,11 +464,10 @@ class TestStoreDecodeRoundTrip: device = torch.device(DEVICE_TYPE) - # Generate rotation - signs = generate_wht_signs(D, seed=42, device=device) + # Pure Hadamard rotation (symmetric: H = H^T, so Pi = PiT = H) H = _build_hadamard(D, DEVICE_TYPE) - PiT = (signs.unsqueeze(1) * H).contiguous().float() - Pi = PiT.T.contiguous() + PiT = H + Pi = H # Generate centroids centroids, _ = solve_lloyd_max(D, cfg.centroid_bits) diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 11985d2fa43..b2a2295ce46 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -406,33 +406,16 @@ class Attention(nn.Module, AttentionLayerBase): def _init_turboquant_buffers( self, cache_dtype: str, head_size: int, prefix: str ) -> None: - """Initialize TurboQuant rotation/projection matrices and centroids.""" + """Initialize TurboQuant centroids for Lloyd-Max quantization.""" from vllm.model_executor.layers.quantization.turboquant.centroids import ( get_centroids, ) from vllm.model_executor.layers.quantization.turboquant.config import ( TurboQuantConfig, ) - from vllm.model_executor.layers.quantization.turboquant.quantizer import ( - generate_wht_signs, - ) tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype, head_size) - # Each layer needs a unique rotation matrix so quantization errors - # don't correlate across layers. Stride must exceed max head_dim to - # ensure non-overlapping RNG streams between adjacent layers. - _TQ_LAYER_SEED_STRIDE = 1337 - - from vllm.model_executor.models.utils import extract_layer_index - - layer_idx = extract_layer_index(prefix) - seed = tq_config.seed + layer_idx * _TQ_LAYER_SEED_STRIDE - - self.register_buffer( - "_tq_signs", - generate_wht_signs(head_size, seed=seed), - ) self.register_buffer( "_tq_centroids", get_centroids(head_size, tq_config.centroid_bits), diff --git a/vllm/model_executor/layers/quantization/turboquant/__init__.py b/vllm/model_executor/layers/quantization/turboquant/__init__.py index 10ee032c9ec..333aef5f825 100644 --- a/vllm/model_executor/layers/quantization/turboquant/__init__.py +++ b/vllm/model_executor/layers/quantization/turboquant/__init__.py @@ -1,12 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""TurboQuant: Near-optimal KV-cache quantization for vLLM. +"""TurboQuant: KV-cache quantization for vLLM. -PolarQuant compression: random rotation + per-coordinate Lloyd-Max -scalar quantization for keys, uniform quantization for values. +Hadamard rotation + per-coordinate Lloyd-Max scalar quantization for +keys, uniform quantization for values. -Reference: "TurboQuant: Online Vector Quantization with Near-optimal -Distortion Rate" (ICLR 2026), Zandieh et al. +The technique implemented here consists of the scalar case of the HIGGS +quantization method (Malinovskii et al., "Pushing the Limits of Large +Language Model Quantization via the Linearity Theorem", NAACL 2025; +preprint arXiv:2411.17525): rotation + optimized grid + optional +re-normalization, applied to KV cache compression. A first application +of this approach to KV-cache compression is in "Cache Me If You Must: +Adaptive Key-Value Quantization for Large Language Models" (Shutova +et al., ICML 2025; preprint arXiv:2501.19392). Both these references +pre-date the TurboQuant paper (Zandieh et al., ICLR 2026). """ from vllm.model_executor.layers.quantization.turboquant.config import TurboQuantConfig diff --git a/vllm/model_executor/layers/quantization/turboquant/config.py b/vllm/model_executor/layers/quantization/turboquant/config.py index 289bed12077..f9cfc89c0c1 100644 --- a/vllm/model_executor/layers/quantization/turboquant/config.py +++ b/vllm/model_executor/layers/quantization/turboquant/config.py @@ -36,10 +36,22 @@ TQ_PRESETS: dict[str, dict] = { class TurboQuantConfig: """Configuration for TurboQuant KV-cache quantization. - Uses PolarQuant (WHT rotation + Lloyd-Max scalar quantization) for keys - and uniform quantization for values. QJL is intentionally omitted — - community consensus (5+ independent groups) found it hurts attention - quality by amplifying variance through softmax. + Applies Hadamard rotation followed by per-coordinate Lloyd-Max scalar + quantization for keys, and uniform quantization for values. + + Historical note: this is the scalar case of the HIGGS quantization + method (Malinovskii et al., "Pushing the Limits of Large Language Model + Quantization via the Linearity Theorem", NAACL 2025; preprint + arXiv:2411.17525): rotation + optimized grid + optional re-normalization, + applied to KV cache compression. A first application of this approach to + KV-cache compression is in "Cache Me If You Must: Adaptive Key-Value + Quantization for Large Language Models" (Shutova et al., ICML 2025; + preprint arXiv:2501.19392). Both these references pre-date the + TurboQuant paper. + + QJL is intentionally omitted — community consensus (5+ independent + groups) found it hurts attention quality by amplifying variance through + softmax. Named presets (use via --kv-cache-dtype): turboquant_k8v4: FP8 keys + 4-bit values, 2.6x, +1.17% PPL @@ -53,8 +65,6 @@ class TurboQuantConfig: rotation/MSE). 3-4 = Lloyd-Max MSE quantized keys. value_quant_bits: Bits per value dimension for uniform quantization. 3 = 8 levels, 4 = 16 levels (default). - seed: Base seed for deterministic random matrix generation. - Actual seed per layer = seed + layer_idx * 1337. norm_correction: Re-normalize centroid vectors to unit norm before inverse rotation during dequant. Fixes quantization-induced norm distortion, improving PPL by ~0.8% at 4-bit. @@ -63,7 +73,7 @@ class TurboQuantConfig: head_dim: int = 128 key_quant_bits: int = 3 # 3-4 = MSE keys, 8 = FP8 keys value_quant_bits: int = 4 # 3-4 = uniform quantized values - seed: int = 42 + seed: int = 42 # kept for backward compatibility; no longer used internally norm_correction: bool = False @property diff --git a/vllm/model_executor/layers/quantization/turboquant/quantizer.py b/vllm/model_executor/layers/quantization/turboquant/quantizer.py index aea63c52bac..82a0c3391ce 100644 --- a/vllm/model_executor/layers/quantization/turboquant/quantizer.py +++ b/vllm/model_executor/layers/quantization/turboquant/quantizer.py @@ -2,23 +2,5 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """TurboQuant quantizer utilities. -Serving path uses generate_wht_signs() for WHT rotation sign buffers. Triton kernels handle all quantization, packing, and dequantization on GPU. """ - -import torch - -_CPU = torch.device("cpu") - - -def generate_wht_signs(d: int, seed: int, device: torch.device = _CPU) -> torch.Tensor: - """Generate deterministic random ±1 signs for WHT rotation. - - Used with Walsh-Hadamard Transform for per-layer rotation randomization. - Same seed derivation as QR (per-layer via seed + layer_idx * stride). - """ - gen = torch.Generator(device="cpu") - gen.manual_seed(seed) - bits = torch.randint(0, 2, (d,), generator=gen, device="cpu") - signs = bits.float() * 2 - 1 - return signs.to(device) diff --git a/vllm/v1/attention/backends/turboquant_attn.py b/vllm/v1/attention/backends/turboquant_attn.py index 2659cbffa82..e7baff9d899 100644 --- a/vllm/v1/attention/backends/turboquant_attn.py +++ b/vllm/v1/attention/backends/turboquant_attn.py @@ -279,29 +279,25 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]): ) def _ensure_on_device(self, layer, device): - """One-time derivation of TQ buffers (rotation matrices, midpoints). + """One-time derivation of TQ buffers (rotation matrix, midpoints). - Registered buffers (_tq_signs, _tq_centroids) are already on the - correct device via register_buffer + model.to(device). + The Hadamard rotation is shared across all layers: random sign + flips do not improve Lloyd-Max quantization quality because the + quantizer is symmetric around zero (sign-flipping a coordinate + maps it to the mirror centroid with identical distortion). """ if not hasattr(layer, "_tq_cached"): - D = layer._tq_signs.shape[0] - signs = layer._tq_signs.to(device=device, dtype=torch.float32) + D = self.head_size - # WHT rotation: orthonormal + self-inverse, enabling future + # Pure Hadamard: orthonormal + symmetric (H = H^T), enabling # in-kernel butterfly fusion and trivial inverse for continuation. H = _build_hadamard(D, str(device)) - layer._tq_PiT = (signs.unsqueeze(1) * H).contiguous() - layer._tq_Pi = layer._tq_PiT.T.contiguous() + layer._tq_PiT = H + layer._tq_Pi = H c = layer._tq_centroids.to(device=device, dtype=torch.float32) - # Precompute midpoints for threshold-based quantization c_sorted, _ = c.sort() layer._tq_midpoints = (c_sorted[:-1] + c_sorted[1:]) / 2 - # Decode buffers (_tq_mid_o_buf, _tq_output_buf, _tq_lse_buf) - # are pre-allocated via register_buffer in Attention.__init__ - # and moved to GPU by model.to(device) — no allocation needed - # here. The memory profiler sees them before KV cache sizing. layer._tq_cached = True def do_kv_cache_update( From d0359f3e0401f3ccb3b8fb66658908c08d265a44 Mon Sep 17 00:00:00 2001 From: mysterious hhhh <69612673+ultranationalism@users.noreply.github.com> Date: Sun, 19 Apr 2026 04:58:46 +0800 Subject: [PATCH 112/150] [Bugfix] Guard mxfp4_experts_quant bindings on ENABLE_NVFP4_SM100 (#40191) Signed-off-by: ultranationalism Signed-off-by: mgoin Co-authored-by: mgoin Co-authored-by: Claude Opus 4.7 (1M context) --- csrc/libtorch_stable/ops.h | 14 -------------- .../quantization/fp4/mxfp4_experts_quant.cu | 10 ++++++++++ csrc/libtorch_stable/torch_bindings.cpp | 8 ++------ 3 files changed, 12 insertions(+), 20 deletions(-) diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index fdf628d6fd9..176cd500633 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -134,20 +134,6 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, torch::stable::Tensor& input, torch::stable::Tensor& input_global_scale); -void mxfp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts, - int64_t n_experts); - -void silu_and_mul_mxfp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts, - int64_t n_experts); - void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, const torch::stable::Tensor& a, const torch::stable::Tensor& b, diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu index 9119f28a750..78e4eda0c01 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu @@ -23,6 +23,7 @@ #include #include +#include #include #include "libtorch_stable/torch_utils.h" #include "libtorch_stable/dispatch_utils.h" @@ -420,3 +421,12 @@ void silu_and_mul_mxfp4_experts_quant( stream); }); } + +// Registered here (not torch_bindings.cpp) because VLLM_GPU_FLAGS is applied +// only under COMPILE_LANGUAGE:CUDA, so ENABLE_NVFP4_SM100 is invisible to +// .cpp files and cannot gate the registration from there. +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant)); + m.impl("silu_and_mul_mxfp4_experts_quant", + TORCH_BOX(&silu_and_mul_mxfp4_experts_quant)); +} diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 95ed6b44f10..124512e8162 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -252,12 +252,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("silu_and_mul_scaled_fp4_experts_quant", TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant)); ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant)); - ops.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant)); - ops.impl("silu_and_mul_mxfp4_experts_quant", - TORCH_BOX(&silu_and_mul_mxfp4_experts_quant)); - - // W4A8 ops: impl registrations are in the source files - // (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu) + // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only). + // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu. #endif } From 38907e4391c29465745f0a9e11dd7e3bae4b30a0 Mon Sep 17 00:00:00 2001 From: Luciano Martins Date: Sat, 18 Apr 2026 20:46:07 -0300 Subject: [PATCH 113/150] [Frontend] Preserve structured output special tokens in offline LLM.chat (#39352) Signed-off-by: Luciano Martins Co-authored-by: Luciano Martins Co-authored-by: Flora Feng <4florafeng@gmail.com> --- vllm/entrypoints/llm.py | 58 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index d135a9ec976..aaef7528253 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1638,6 +1638,17 @@ class LLM: seq_params = self._params_to_seq(params, len(seq_convs)) seq_lora_requests = self._lora_request_to_seq(lora_request, len(seq_convs)) + # When thinking is enabled or tools are provided, and the model + # uses special tokens for structured output (e.g. Gemma4's + # <|channel>, <|tool_call>, <|"|>), automatically set + # skip_special_tokens=False so these tokens are preserved in + # output.text for downstream parsing. + needs_parsing = ( + chat_template_kwargs and chat_template_kwargs.get("enable_thinking") + ) or tools + if needs_parsing: + self._adjust_params_for_parsing(seq_params) + return self._render_and_run_requests( prompts=( self._preprocess_chat_one( @@ -1663,6 +1674,53 @@ class LLM: use_tqdm=use_tqdm, ) + def _adjust_params_for_parsing( + self, params: Sequence[SamplingParams | PoolingParams] + ) -> None: + """Set ``skip_special_tokens=False`` when the model encodes + structured output syntax as special tokens. + + Models like Gemma4 register thinking delimiters + (``<|channel>``/````) and tool call tokens + (``<|tool_call>``/````/``<|"|>``) as special tokens. + The default ``skip_special_tokens=True`` strips them from + ``output.text``, breaking parsing of both reasoning blocks and + tool calls. + + This is a no-op for models whose structured tokens are regular + text tokens (e.g. DeepSeek's ````/````). + """ + # The offline API currently lacks a unified rendering pipeline. + # Until the planned Renderer refactor is complete, we hardcode + # this token preservation logic specifically for Gemma4 models + # to avoid regressions on other models. + hf_config = getattr(self.model_config, "hf_config", None) + architectures = getattr(hf_config, "architectures", []) + + if any("Gemma4" in arch for arch in architectures): + tokenizer = self.renderer.get_tokenizer() + vocab = tokenizer.get_vocab() + special_ids = set(getattr(tokenizer, "all_special_ids", [])) + + # Tokens used for thinking delimiters and tool call syntax + # that some models (Gemma4) register as special tokens. + structured_tokens = ( + "<|channel>", + "", # thinking delimiters + "<|tool_call>", + "", # tool call delimiters + '<|"|>', # string quoting in tool args + ) + needs_special = any( + vocab.get(tok) in special_ids + for tok in structured_tokens + if tok in vocab + ) + if needs_special: + for sp in params: + if isinstance(sp, SamplingParams) and sp.skip_special_tokens: + sp.skip_special_tokens = False + def _render_and_run_requests( self, prompts: Iterable[EngineInput], From 4b7f5ea1a06c08542af6dedbc9086750d3e2666a Mon Sep 17 00:00:00 2001 From: omerpaz95 <73347585+omerpaz95@users.noreply.github.com> Date: Sun, 19 Apr 2026 07:49:10 +0300 Subject: [PATCH 114/150] [KV Connector] Allow metrics of multiple connectors of same types in multi connector. (#40010) Signed-off-by: omerpaz95 Co-authored-by: Or Ozeri --- .../kv_transfer/kv_connector/v1/multi_connector.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 3888d2e0f44..4ef8f0ac9c9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -548,7 +548,13 @@ class MultiConnector(KVConnectorBase_V1): if stats_by_connector is None: # Lazy init to allow optional return value. stats_by_connector = MultiKVConnectorStats() - stats_by_connector[c.__class__.__name__] = stats + connector_id = c.__class__.__name__ + if connector_id in stats_by_connector.data: + stats_by_connector[connector_id] = stats_by_connector[ + connector_id + ].aggregate(stats) + else: + stats_by_connector[connector_id] = stats return stats_by_connector @classmethod @@ -560,9 +566,13 @@ class MultiConnector(KVConnectorBase_V1): per_engine_labelvalues: dict[int, list[object]], ) -> KVConnectorPromMetrics: prom_metrics: dict[str, KVConnectorPromMetrics] = {} + seen_classes: set[type] = set() for connector_cls, temp_config in cls._get_connector_classes_and_configs( vllm_config ): + if connector_cls in seen_classes: + continue + seen_classes.add(connector_cls) connector_prom = connector_cls.build_prom_metrics( temp_config, metric_types, labelnames, per_engine_labelvalues ) From 4353c9cb4ac8d9818303c2e61282062aff4074c3 Mon Sep 17 00:00:00 2001 From: omerpaz95 <73347585+omerpaz95@users.noreply.github.com> Date: Sun, 19 Apr 2026 08:54:59 +0300 Subject: [PATCH 115/150] [KV Offload] Pass request context (#39185) Signed-off-by: omerpaz95 --- .../offloading_connector/test_scheduler.py | 54 +++++--- .../unit/offloading_connector/utils.py | 2 +- tests/v1/kv_offload/test_cpu_manager.py | 117 ++++++++++-------- .../kv_connector/v1/offloading/scheduler.py | 14 ++- vllm/v1/kv_offload/abstract.py | 28 ++++- vllm/v1/kv_offload/cpu/manager.py | 19 ++- vllm/v1/kv_offload/reuse_manager.py | 17 ++- 7 files changed, 170 insertions(+), 81 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index bdc81dc1ac3..26bd01b138f 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -32,8 +32,8 @@ def test_offloading_connector(request_runner, async_scheduling: bool): # 3 blocks, store just the middle block (skip first and last) # blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8] runner.new_request(token_ids=[0] * offloaded_block_size * 3) - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output( - list(keys)[1:2] + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output(list(keys)[1:2]) ) runner.run(decoded_tokens=[0]) @@ -45,18 +45,22 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.manager.prepare_store.assert_not_called() # +1 token -> single block, fail prepare_store - runner.manager.prepare_store.side_effect = lambda keys: None + runner.manager.prepare_store.side_effect = lambda keys, req_context: None runner.run(decoded_tokens=[0]) runner.manager.prepare_store.assert_called() # 1 more block (+ token for async scheduling) # now set block_hashes_to_store = [] - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([]) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output([]) + ) runner.run(decoded_tokens=[0] * (offloaded_block_size + 1)) # 1 more block (+ token for kicking off offloading) # now check touch was called with all 6 blocks - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output(keys) + ) runner.run( decoded_tokens=[0] * (offloaded_block_size + 1), expected_stored_gpu_block_indexes=(15, 16, 17), @@ -89,13 +93,17 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.new_request( token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size) ) - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([]) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output([]) + ) runner.run(decoded_tokens=[EOS_TOKEN_ID]) runner.manager.lookup.assert_not_called() # single block lookup with no hits runner.new_request(token_ids=[1] * offloaded_block_size) - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([]) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output([]) + ) runner.run(decoded_tokens=[EOS_TOKEN_ID]) runner.manager.lookup.assert_called() assert len(list(runner.manager.lookup.call_args.args[0])) == 1 @@ -103,7 +111,9 @@ def test_offloading_connector(request_runner, async_scheduling: bool): # single block lookup with a hit runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * offloaded_block_size) - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([]) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output([]) + ) runner.manager.lookup.return_value = 1 runner.run( decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2) @@ -113,7 +123,9 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.new_request( token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size ) - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([]) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output([]) + ) runner.manager.lookup.return_value = 1 runner.run( decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5) @@ -164,14 +176,18 @@ def test_request_preemption(request_runner, async_scheduling: bool): # 2 blocks, store all, without flushing # blocks = [0, 1, 2], [3, 4, 5] runner.new_request(token_ids=[0] * offloaded_block_size * 2) - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output(keys) + ) runner.run( decoded_tokens=[0], complete_transfers=False, ) # decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush) - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output(keys) + ) runner.run( decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size), complete_transfers=False, @@ -195,7 +211,9 @@ def test_request_preemption(request_runner, async_scheduling: bool): # request should now return from preemption # re-load [0, ..., 8] from the CPU and store [9, 10, 11] runner.manager.lookup.return_value = 3 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output(keys) + ) runner.run( decoded_tokens=[0] * gpu_block_size, expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8), @@ -222,7 +240,9 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: # store 1 blocks runner.new_request(token_ids=[0] * offloaded_block_size) - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output(keys) + ) runner.run( decoded_tokens=[EOS_TOKEN_ID], expected_stored_gpu_block_indexes=(0, 1, 2), @@ -253,7 +273,9 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs) # complete transfers - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([]) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output([]) + ) runner.run( decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2), @@ -278,7 +300,9 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool): # store 1 blocks runner.new_request(token_ids=[0] * offloaded_block_size) - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output(keys) + ) runner.run( decoded_tokens=[EOS_TOKEN_ID], expected_stored_gpu_block_indexes=(0, 1, 2), diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index d774d81def9..aaf9152a43c 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -115,7 +115,7 @@ class MockOffloadingSpec(OffloadingSpec): self.manager = MagicMock(spec=OffloadingManager) self.manager.lookup.return_value = 0 - self.manager.prepare_load = lambda keys: MockLoadStoreSpec(keys) + self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys) self.handler = MockOffloadingHandler() def get_manager(self) -> OffloadingManager: diff --git a/tests/v1/kv_offload/test_cpu_manager.py b/tests/v1/kv_offload/test_cpu_manager.py index 7a2ba837eca..651ea091ae6 100644 --- a/tests/v1/kv_offload/test_cpu_manager.py +++ b/tests/v1/kv_offload/test_cpu_manager.py @@ -11,6 +11,7 @@ from vllm.v1.kv_offload.abstract import ( OffloadingEvent, OffloadKey, PrepareStoreOutput, + ReqContext, make_offload_key, ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager @@ -19,6 +20,14 @@ from vllm.v1.kv_offload.mediums import CPULoadStoreSpec from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager +def make_req_context(kv_transfer_params: dict | None = None) -> ReqContext: + """Create a ReqContext as production code would, from a request's params.""" + return ReqContext(kv_transfer_params=kv_transfer_params) + + +_EMPTY_REQ_CTX = make_req_context() + + @dataclass class ExpectedPrepareStoreOutput: keys_to_store: list[int] @@ -103,7 +112,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): ) # store [1, 2] and complete - manager.prepare_store(to_keys([1, 2])) + manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) manager.complete_store(to_keys([1, 2])) # touch [1] to make block 2 the LRU candidate @@ -113,7 +122,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): # - block 2 is already stored -> filtered out of keys_to_store # - block 2 must NOT be evicted even though it is the LRU candidate # - block 1 (ID 0) is evicted instead; new blocks [3,4,5] get IDs 2,3,0 - prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5])) + prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( @@ -127,7 +136,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): manager.complete_store(to_keys([2, 3, 4, 5])) # block 2 must still be present in the cache - assert manager.lookup(to_keys([2])) == 1 + assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1 def test_cpu_manager(): @@ -140,7 +149,7 @@ def test_cpu_manager(): ) # prepare store [1, 2] - prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2])) + prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( @@ -151,7 +160,7 @@ def test_cpu_manager(): ) # lookup [1, 2] -> not ready - assert cpu_manager.lookup(to_keys([1, 2])) == 0 + assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 # no events so far assert list(cpu_manager.take_events()) == [] @@ -161,12 +170,14 @@ def test_cpu_manager(): verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] - assert cpu_manager.lookup(to_keys([1])) == 1 - assert cpu_manager.lookup(to_keys([1, 2])) == 2 - assert cpu_manager.lookup(to_keys([1, 2, 3])) == 2 + assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1 + assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2 + assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2 # prepare store [2, 3, 4, 5] -> evicts [1] - prepare_store_output = cpu_manager.prepare_store(to_keys([2, 3, 4, 5])) + prepare_store_output = cpu_manager.prepare_store( + to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX + ) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( @@ -180,23 +191,23 @@ def test_cpu_manager(): verify_events(cpu_manager.take_events(), expected_evictions=({1},)) # prepare store with no space - assert cpu_manager.prepare_store(to_keys([1, 6])) is None + assert cpu_manager.prepare_store(to_keys([1, 6]), _EMPTY_REQ_CTX) is None # complete store [2, 3, 4, 5] cpu_manager.complete_store(to_keys([2, 3, 4, 5])) # prepare load [2, 3] - prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3])) + prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX) verify_load_output(prepare_load_output, [1, 2]) # prepare store with no space ([2, 3] is being loaded) - assert cpu_manager.prepare_store(to_keys([6, 7, 8])) is None + assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None # complete load [2, 3] cpu_manager.complete_load(to_keys([2, 3])) # prepare store [6, 7, 8] -> evicts [2, 3, 4] (oldest) - prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8])) + prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( @@ -213,7 +224,7 @@ def test_cpu_manager(): cpu_manager.touch(to_keys([5, 6, 7])) # prepare store [7, 9] -> evicts [8] (oldest following previous touch) - prepare_store_output = cpu_manager.prepare_store(to_keys([9])) + prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( @@ -227,8 +238,8 @@ def test_cpu_manager(): cpu_manager.complete_store(to_keys([7, 9]), success=False) # assert [7] is still stored, but [9] is not - assert cpu_manager.lookup(to_keys([7])) == 1 - assert cpu_manager.lookup(to_keys([9])) == 0 + assert cpu_manager.lookup(to_keys([7]), _EMPTY_REQ_CTX) == 1 + assert cpu_manager.lookup(to_keys([9]), _EMPTY_REQ_CTX) == 0 verify_events( cpu_manager.take_events(), @@ -260,7 +271,9 @@ class TestARCPolicy: cpu_manager, arc_policy = self._make_manager() # prepare store [1, 2] - prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2])) + prepare_store_output = cpu_manager.prepare_store( + to_keys([1, 2]), _EMPTY_REQ_CTX + ) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( @@ -271,7 +284,7 @@ class TestARCPolicy: ) # lookup [1, 2] -> not ready - assert cpu_manager.lookup(to_keys([1, 2])) == 0 + assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 # no events so far assert list(cpu_manager.take_events()) == [] @@ -281,9 +294,9 @@ class TestARCPolicy: verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] - assert cpu_manager.lookup(to_keys([1])) == 1 - assert cpu_manager.lookup(to_keys([1, 2])) == 2 - assert cpu_manager.lookup(to_keys([1, 2, 3])) == 2 + assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1 + assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2 + assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2 # blocks should be in T1 (recent) assert len(arc_policy.t1) == 2 @@ -297,7 +310,7 @@ class TestARCPolicy: cpu_manager, arc_policy = self._make_manager(enable_events=False) # store and complete block 1 - cpu_manager.prepare_store(to_keys([1])) + cpu_manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1])) # block 1 starts in T1 (recent) @@ -319,7 +332,9 @@ class TestARCPolicy: cpu_manager, _ = self._make_manager() # prepare and complete store [1, 2, 3, 4] - prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2, 3, 4])) + prepare_store_output = cpu_manager.prepare_store( + to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX + ) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( @@ -331,19 +346,21 @@ class TestARCPolicy: cpu_manager.complete_store(to_keys([1, 2, 3, 4])) # prepare load [2, 3] (increases ref_cnt) - prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3])) + prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX) verify_load_output(prepare_load_output, [1, 2]) # prepare store [5, 6, 7] with [2, 3] being loaded # should fail because [2, 3] have ref_cnt > 0 - assert cpu_manager.prepare_store(to_keys([5, 6, 7])) is None + assert cpu_manager.prepare_store(to_keys([5, 6, 7]), _EMPTY_REQ_CTX) is None # complete load [2, 3] cpu_manager.complete_load(to_keys([2, 3])) # now prepare store [5, 6, 7] should succeed # ARC will evict blocks one at a time from T1 as needed - prepare_store_output = cpu_manager.prepare_store(to_keys([5, 6, 7])) + prepare_store_output = cpu_manager.prepare_store( + to_keys([5, 6, 7]), _EMPTY_REQ_CTX + ) assert prepare_store_output is not None # Should successfully evict enough blocks to make room (at least 1) assert len(prepare_store_output.evicted_keys) >= 1 @@ -357,13 +374,13 @@ class TestARCPolicy: cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False) # store blocks 1, 2 (fills cache) - cpu_manager.prepare_store(to_keys([1, 2])) + cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1, 2])) initial_target = arc_policy.target_t1_size # store block 3, evicting block 1 (moves to B1 ghost list) - cpu_manager.prepare_store(to_keys([3])) + cpu_manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([3])) # block 1 should be in B1 (ghost list) @@ -384,7 +401,7 @@ class TestARCPolicy: cpu_manager, arc_policy = self._make_manager(enable_events=False) # store blocks 1, 2, 3, 4 - cpu_manager.prepare_store(to_keys([1, 2, 3, 4])) + cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1, 2, 3, 4])) # promote blocks 3, 4 to T2 by touching them @@ -399,7 +416,7 @@ class TestARCPolicy: arc_policy.target_t1_size = 1 # store block 5, should evict from T1 (block 1, LRU in T1) - output = cpu_manager.prepare_store(to_keys([5])) + output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX) assert output is not None assert to_keys([1]) == output.evicted_keys @@ -418,12 +435,12 @@ class TestARCPolicy: cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False) # fill cache with blocks 1, 2 - cpu_manager.prepare_store(to_keys([1, 2])) + cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1, 2])) # store many blocks to fill ghost lists for i in range(3, 20): - cpu_manager.prepare_store(to_keys([i])) + cpu_manager.prepare_store(to_keys([i]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([i])) # ghost lists should not exceed cache_capacity @@ -438,7 +455,7 @@ class TestARCPolicy: cpu_manager, arc_policy = self._make_manager() # store blocks 1, 2, 3, 4 - cpu_manager.prepare_store(to_keys([1, 2, 3, 4])) + cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1, 2, 3, 4])) # promote 3, 4 to T2 @@ -453,7 +470,7 @@ class TestARCPolicy: assert len(arc_policy.t2) == 3 # store block 5, should evict from T1 (block 2, only one in T1) - prepare_store_output = cpu_manager.prepare_store(to_keys([5])) + prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( @@ -471,11 +488,11 @@ class TestARCPolicy: cpu_manager, arc_policy = self._make_manager() # store blocks 1, 2, 3, 4 - cpu_manager.prepare_store(to_keys([1, 2, 3, 4])) + cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1, 2, 3, 4])) # prepare store block 5 (will evict block 1) - prepare_store_output = cpu_manager.prepare_store(to_keys([5])) + prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX) assert prepare_store_output is not None assert len(prepare_store_output.evicted_keys) == 1 @@ -483,7 +500,7 @@ class TestARCPolicy: cpu_manager.complete_store(to_keys([5]), success=False) # block 5 should not be in cache - assert cpu_manager.lookup(to_keys([5])) == 0 + assert cpu_manager.lookup(to_keys([5]), _EMPTY_REQ_CTX) == 0 # block 5 should not be in T1 or T2 assert to_keys([5])[0] not in arc_policy.t1 assert to_keys([5])[0] not in arc_policy.t2 @@ -500,11 +517,13 @@ class TestARCPolicy: cpu_manager, arc_policy = self._make_manager() # store [1, 2] - cpu_manager.prepare_store(to_keys([1, 2])) + cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1, 2])) # store [3, 4, 5] -> evicts [1] - prepare_store_output = cpu_manager.prepare_store(to_keys([3, 4, 5])) + prepare_store_output = cpu_manager.prepare_store( + to_keys([3, 4, 5]), _EMPTY_REQ_CTX + ) assert prepare_store_output is not None assert len(prepare_store_output.evicted_keys) == 1 cpu_manager.complete_store(to_keys([3, 4, 5])) @@ -517,13 +536,13 @@ class TestARCPolicy: assert len(arc_policy.t2) == 2 # store [6] -> should evict from T1 (4 is oldest in T1) - prepare_store_output = cpu_manager.prepare_store(to_keys([6])) + prepare_store_output = cpu_manager.prepare_store(to_keys([6]), _EMPTY_REQ_CTX) assert prepare_store_output is not None cpu_manager.complete_store(to_keys([6])) # verify blocks 2, 3 (in T2) are still present - assert cpu_manager.lookup(to_keys([2])) == 1 - assert cpu_manager.lookup(to_keys([3])) == 1 + assert cpu_manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1 + assert cpu_manager.lookup(to_keys([3]), _EMPTY_REQ_CTX) == 1 # verify events events = list(cpu_manager.take_events()) @@ -543,34 +562,34 @@ def test_filter_reused_manager(): ) # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet - assert manager.lookup(to_keys([1, 2])) == 0 + assert manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 # prepare store [1, 2] -> should be filtered - prepare_store_output = manager.prepare_store(to_keys([1, 2])) + prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) assert prepare_store_output is not None assert prepare_store_output.keys_to_store == [] # Lookup [1] -> 2nd time, eligible now - assert manager.lookup(to_keys([1])) == 0 + assert manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 0 # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered - prepare_store_output = manager.prepare_store(to_keys([1, 2])) + prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) assert prepare_store_output is not None assert prepare_store_output.keys_to_store == to_keys([1]) # Lookup [3, 4] -> 1st time # (evicts [2] from tracker since max_size is 3 and tracker has [1]) - assert manager.lookup(to_keys([3, 4])) == 0 + assert manager.lookup(to_keys([3, 4]), _EMPTY_REQ_CTX) == 0 # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4]) assert to_keys([2])[0] not in manager.counts # Lookup [2] again -> (this adds [2] back to the tracker as 1st time) - assert manager.lookup(to_keys([2])) == 0 + assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 0 # Verify [2] was re-added with count=1 (not eligible yet) assert manager.counts.get(to_keys([2])[0]) == 1 # prepare store [2] -> should still be filtered out since count was reset - prepare_store_output = manager.prepare_store(to_keys([2])) + prepare_store_output = manager.prepare_store(to_keys([2]), _EMPTY_REQ_CTX) assert prepare_store_output is not None assert prepare_store_output.keys_to_store == [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 9fd0bed8d3e..c5272ea2778 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -19,6 +19,7 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_offload.abstract import ( OffloadingManager, OffloadKey, + ReqContext, get_offload_block_hash, make_offload_key, ) @@ -74,6 +75,7 @@ class RequestOffloadState: config: SchedulerOffloadConfig req: Request group_states: tuple[RequestGroupState, ...] = field(init=False) + req_context: ReqContext = field(init=False) # number of hits in the GPU cache num_locally_computed_tokens: int = 0 @@ -81,6 +83,7 @@ class RequestOffloadState: self.group_states = tuple( RequestGroupState() for _ in self.config.kv_group_configs ) + self.req_context = ReqContext(kv_transfer_params=self.req.kv_transfer_params) def update_offload_keys(self) -> None: for group_config, group_state in zip( @@ -181,7 +184,10 @@ class OffloadingConnectorScheduler: return 0, False start_block_idx = num_computed_tokens // group_config.offloaded_block_size - hits = self.manager.lookup(offload_keys[start_block_idx:]) + hits = self.manager.lookup( + offload_keys[start_block_idx:], + req_status.req_context, + ) if hits is None: # indicates a lookup that should be tried later return None, False @@ -249,7 +255,7 @@ class OffloadingConnectorScheduler: assert len(request.block_hashes) // self.config.block_size_factor >= num_blocks offload_keys = group_state.offload_keys[start_block_idx:num_blocks] - src_spec = self.manager.prepare_load(offload_keys) + src_spec = self.manager.prepare_load(offload_keys, req_status.req_context) dst_spec = GPULoadStoreSpec( block_ids[num_computed_gpu_blocks:], group_sizes=(num_pending_gpu_blocks,), @@ -304,7 +310,9 @@ class OffloadingConnectorScheduler: assert len(req.block_hashes) >= num_gpu_blocks new_offload_keys = group_state.offload_keys[start_block_idx:num_blocks] - store_output = self.manager.prepare_store(new_offload_keys) + store_output = self.manager.prepare_store( + new_offload_keys, req_status.req_context + ) if store_output is None: logger.warning( "Request %s: cannot store %s blocks", req_id, num_new_blocks diff --git a/vllm/v1/kv_offload/abstract.py b/vllm/v1/kv_offload/abstract.py index efa617b80f9..39ee360e8f6 100644 --- a/vllm/v1/kv_offload/abstract.py +++ b/vllm/v1/kv_offload/abstract.py @@ -30,7 +30,7 @@ The class provides the following primitives: from abc import ABC, abstractmethod from collections.abc import Iterable from dataclasses import dataclass -from typing import NewType +from typing import Any, NewType # `OffloadKey` identifies an offloaded block. It combines a block hash with # its KV cache group index, encoded as raw bytes to avoid tuple GC overhead. @@ -53,6 +53,11 @@ def get_offload_group_idx(key: OffloadKey) -> int: return int.from_bytes(key[-4:], "big", signed=False) +@dataclass +class ReqContext: + kv_transfer_params: dict[str, Any] | None = None + + class LoadStoreSpec(ABC): """ Abstract metadata that encapsulates information allowing a worker @@ -86,13 +91,18 @@ class OffloadingEvent: class OffloadingManager(ABC): @abstractmethod - def lookup(self, keys: Iterable[OffloadKey]) -> int | None: + def lookup( + self, + keys: Iterable[OffloadKey], + req_context: ReqContext, + ) -> int | None: """ Finds the length of the maximal series of blocks, starting from the first one, that are all offloaded. Args: keys: the keys identifying the blocks to lookup. + req_context: per-request context (e.g. kv_transfer_params). Returns: An integer representing the maximal number of blocks that @@ -103,7 +113,11 @@ class OffloadingManager(ABC): pass @abstractmethod - def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec: + def prepare_load( + self, + keys: Iterable[OffloadKey], + req_context: ReqContext, + ) -> LoadStoreSpec: """ Prepare the given blocks to be read. The given blocks will be protected from eviction until @@ -112,6 +126,7 @@ class OffloadingManager(ABC): Args: keys: the keys identifying the blocks. + req_context: per-request context (e.g. kv_transfer_params). Returns: A LoadStoreSpec that can be used by a worker to locate and load @@ -139,7 +154,11 @@ class OffloadingManager(ABC): return @abstractmethod - def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None: + def prepare_store( + self, + keys: Iterable[OffloadKey], + req_context: ReqContext, + ) -> PrepareStoreOutput | None: """ Prepare the given blocks to be offloaded. The given blocks will be protected from eviction until @@ -147,6 +166,7 @@ class OffloadingManager(ABC): Args: keys: the keys identifying the blocks. + req_context: per-request context (e.g. kv_transfer_params). Returns: A PrepareStoreOutput indicating which blocks need storing, diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 2befa91c77a..5ae7454430f 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -9,6 +9,7 @@ from vllm.v1.kv_offload.abstract import ( OffloadingManager, OffloadKey, PrepareStoreOutput, + ReqContext, ) from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy @@ -83,7 +84,11 @@ class CPUOffloadingManager(OffloadingManager): # --- OffloadingManager interface --- - def lookup(self, keys: Iterable[OffloadKey]) -> int | None: + def lookup( + self, + keys: Iterable[OffloadKey], + req_context: ReqContext, + ) -> int | None: hit_count = 0 for key in keys: block = self._policy.get(key) @@ -92,7 +97,11 @@ class CPUOffloadingManager(OffloadingManager): hit_count += 1 return hit_count - def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec: + def prepare_load( + self, + keys: Iterable[OffloadKey], + req_context: ReqContext, + ) -> LoadStoreSpec: blocks = [] for key in keys: block = self._policy.get(key) @@ -112,7 +121,11 @@ class CPUOffloadingManager(OffloadingManager): assert block.ref_cnt > 0, f"Block {key!r} ref_cnt is already 0" block.ref_cnt -= 1 - def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None: + def prepare_store( + self, + keys: Iterable[OffloadKey], + req_context: ReqContext, + ) -> PrepareStoreOutput | None: keys_list = list(keys) # filter out blocks that are already stored diff --git a/vllm/v1/kv_offload/reuse_manager.py b/vllm/v1/kv_offload/reuse_manager.py index 95de602976a..a9650e38c51 100644 --- a/vllm/v1/kv_offload/reuse_manager.py +++ b/vllm/v1/kv_offload/reuse_manager.py @@ -16,6 +16,7 @@ from vllm.v1.kv_offload.abstract import ( OffloadingManager, OffloadKey, PrepareStoreOutput, + ReqContext, ) @@ -65,7 +66,7 @@ class FilterReusedOffloadingManager(OffloadingManager): # Intercepted methods # ------------------------------------------------------------------ - def lookup(self, keys: Iterable[OffloadKey]) -> int | None: + def lookup(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> int | None: """Record each key, then delegate lookup to backing manager.""" keys = list(keys) for key in keys: @@ -76,9 +77,11 @@ class FilterReusedOffloadingManager(OffloadingManager): if len(self.counts) >= self.max_tracker_size: self.counts.popitem(last=False) # evict LRU self.counts[key] = 1 - return self._backing.lookup(keys) + return self._backing.lookup(keys, req_context) - def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None: + def prepare_store( + self, keys: Iterable[OffloadKey], req_context: ReqContext + ) -> PrepareStoreOutput | None: """Filter out blocks below threshold, then delegate to backing. Filtering is evaluated *before* calling the backing manager's @@ -93,14 +96,16 @@ class FilterReusedOffloadingManager(OffloadingManager): # Passing an empty list is intentional and safe — CPUOffloadingManager # handles it correctly, returning a PrepareStoreOutput with empty lists. # Delegate to the backing manager with only the eligible keys. - return self._backing.prepare_store(eligible) + return self._backing.prepare_store(eligible, req_context) # ------------------------------------------------------------------ # Delegated methods # ------------------------------------------------------------------ - def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec: - return self._backing.prepare_load(keys) + def prepare_load( + self, keys: Iterable[OffloadKey], req_context: ReqContext + ) -> LoadStoreSpec: + return self._backing.prepare_load(keys, req_context) def touch(self, keys: Iterable[OffloadKey]) -> None: return self._backing.touch(keys) From 03ce1c6ed908788cf508219b666720783a723123 Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Sun, 19 Apr 2026 04:30:27 -0400 Subject: [PATCH 116/150] [Bugfix] Kimi-K2 tool parser streaming - fix token leakage, argument truncation, and content dropping (#38579) Signed-off-by: sfeng33 <4florafeng@gmail.com> --- .../tool_parsers/test_kimi_k2_tool_parser.py | 1468 ++++++----------- vllm/tool_parsers/kimi_k2_tool_parser.py | 637 ++----- 2 files changed, 692 insertions(+), 1413 deletions(-) diff --git a/tests/tool_parsers/test_kimi_k2_tool_parser.py b/tests/tool_parsers/test_kimi_k2_tool_parser.py index 09c1c461c10..b56032b91c1 100644 --- a/tests/tool_parsers/test_kimi_k2_tool_parser.py +++ b/tests/tool_parsers/test_kimi_k2_tool_parser.py @@ -3,14 +3,20 @@ # ruff: noqa: E501 import json +from unittest.mock import MagicMock import pytest -from vllm.entrypoints.openai.engine.protocol import FunctionCall, ToolCall +from tests.tool_parsers.utils import ( + run_tool_extraction, + run_tool_extraction_streaming, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) from vllm.tokenizers import get_tokenizer from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser -# Use a common model that is likely to be available MODEL = "moonshotai/Kimi-K2-Instruct" @@ -20,959 +26,557 @@ def kimi_k2_tokenizer(): @pytest.fixture -def kimi_k2_tool_parser(kimi_k2_tokenizer): +def parser(kimi_k2_tokenizer): return KimiK2ToolParser(kimi_k2_tokenizer) -def assert_tool_calls( - actual_tool_calls: list[ToolCall], expected_tool_calls: list[ToolCall] -): - assert len(actual_tool_calls) == len(expected_tool_calls) +SECTION_BEGIN = "<|tool_calls_section_begin|>" +SECTION_END = "<|tool_calls_section_end|>" +TOOL_BEGIN = "<|tool_call_begin|>" +TOOL_END = "<|tool_call_end|>" +ARG_BEGIN = "<|tool_call_argument_begin|>" - for actual_tool_call, expected_tool_call in zip( - actual_tool_calls, expected_tool_calls + +def _tool(tool_id: str, args: str) -> str: + return f"{TOOL_BEGIN}{tool_id} {ARG_BEGIN}{args}{TOOL_END}" + + +def _wrap(*tool_strs: str) -> str: + return SECTION_BEGIN + "".join(tool_strs) + SECTION_END + + +class TestExtractToolCalls: + def test_no_tools(self, parser): + content, tool_calls = run_tool_extraction( + parser, "This is a test", streaming=False + ) + assert content == "This is a test" + assert tool_calls == [] + + @pytest.mark.parametrize( + "model_output, expected_names, expected_args_list, expected_content", + [ + pytest.param( + "I'll check. " + + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')), + ["get_weather"], + [{"city": "Beijing"}], + "I'll check. ", + id="single_tool_call", + ), + pytest.param( + "Compare weather. " + + _wrap( + _tool("functions.get_weather:0", '{"city": "Beijing"}'), + _tool("functions.get_weather:1", '{"city": "Shanghai"}'), + ), + ["get_weather", "get_weather"], + [{"city": "Beijing"}, {"city": "Shanghai"}], + "Compare weather. ", + id="parallel_tool_calls", + ), + pytest.param( + "Multiple tasks. " + + _wrap( + _tool("functions.get_weather:0", '{"city": "New York"}'), + _tool("functions.get_news:1", '{"topic": "technology"}'), + _tool( + "functions.send_email:2", + '{"to": "user@example.com", "subject": "Daily Update"}', + ), + ), + ["get_weather", "get_news", "send_email"], + [ + {"city": "New York"}, + {"topic": "technology"}, + {"to": "user@example.com", "subject": "Daily Update"}, + ], + "Multiple tasks. ", + id="three_tool_calls", + ), + pytest.param( + "Process HTML. " + + _wrap( + _tool("functions.process_html:0", '{"html": "
content
"}') + ), + ["process_html"], + [{"html": "
content
"}], + "Process HTML. ", + id="angle_brackets_in_json", + ), + pytest.param( + "Formatted. " + + _wrap( + _tool( + "functions.process_data:0", + '{\n "name": "test",\n "value": 123\n}', + ) + ), + ["process_data"], + [{"name": "test", "value": 123}], + "Formatted. ", + id="multiline_json", + ), + pytest.param( + "No prefix. " + _wrap(_tool("get_weather:0", '{"city": "Tokyo"}')), + ["get_weather"], + [{"city": "Tokyo"}], + "No prefix. ", + id="no_functions_prefix", + ), + pytest.param( + "Empty args. " + _wrap(_tool("functions.test:0", "{}")), + ["test"], + [{}], + "Empty args. ", + id="empty_arguments", + ), + ], + ) + def test_extract_tool_calls( + self, parser, model_output, expected_names, expected_args_list, expected_content ): - assert actual_tool_call.type == "function" - assert actual_tool_call.function == expected_tool_call.function + content, tool_calls = run_tool_extraction(parser, model_output, streaming=False) + assert content == expected_content + assert len(tool_calls) == len(expected_names) + for tc, name, expected_args in zip( + tool_calls, expected_names, expected_args_list + ): + assert tc.type == "function" + assert tc.function.name == name + assert json.loads(tc.function.arguments) == expected_args + # id format: "something:digit" + assert tc.id.split(":")[-1].isdigit() - # assert tool call id format: should contain function name and numeric index - # Format can be either "functions.func_name:0" or "func_name:0" - assert actual_tool_call.id.split(":")[-1].isdigit() - assert ( - actual_tool_call.id.split(":")[0].split(".")[-1] - == expected_tool_call.function.name + def test_invalid_json_still_extracted(self, parser): + """Tool calls with invalid JSON are still returned (arguments as-is).""" + model_output = ( + "Help. " + + SECTION_BEGIN + + _tool("functions.bad:0", '{"city": "Beijing"') + + _tool("functions.good:1", '{"city": "Shanghai"}') + + SECTION_END + ) + content, tool_calls = run_tool_extraction(parser, model_output, streaming=False) + assert len(tool_calls) == 2 + assert tool_calls[0].function.name == "bad" + assert tool_calls[1].function.name == "good" + + def test_invalid_funcall_id_skipped(self, parser): + """Tool calls with malformed id (no colon+digit) are skipped.""" + model_output = ( + "Help. " + + SECTION_BEGIN + + _tool("functions.invalid.0", '{"city": "Beijing"}') + + _tool("functions.valid:1", '{"city": "Shanghai"}') + + SECTION_END + ) + content, tool_calls = run_tool_extraction(parser, model_output, streaming=False) + assert len(tool_calls) == 1 + assert tool_calls[0].function.name == "valid" + + def test_native_id_extracted(self, parser): + """Regression: parser extracts native ID onto ToolCall (PR #32768).""" + model_output = "Checking weather. " + _wrap( + _tool("functions.get_weather:0", '{"city": "Tokyo"}') + ) + content, tool_calls = run_tool_extraction(parser, model_output, streaming=False) + assert len(tool_calls) == 1 + assert tool_calls[0].id == "functions.get_weather:0" + assert tool_calls[0].function.name == "get_weather" + assert json.loads(tool_calls[0].function.arguments) == {"city": "Tokyo"} + + def test_multi_turn_native_id_continuity(self, kimi_k2_tokenizer): + """Regression: native IDs from turn 1 preserved across turns (PR #32768).""" + turn1_parser = KimiK2ToolParser(kimi_k2_tokenizer) + turn1_output = "Let me check. " + _wrap( + _tool("functions.get_weather:0", '{"city": "Beijing"}') + ) + _, turn1_tools = run_tool_extraction( + turn1_parser, turn1_output, streaming=False + ) + assert len(turn1_tools) == 1 + assert turn1_tools[0].id == "functions.get_weather:0" + + # Fresh parser for turn 2 + turn2_parser = KimiK2ToolParser(kimi_k2_tokenizer) + turn2_output = "Now let me get news. " + _wrap( + _tool("functions.get_news:0", '{"topic": "weather in Beijing"}') + ) + _, turn2_tools = run_tool_extraction( + turn2_parser, turn2_output, streaming=False + ) + assert len(turn2_tools) == 1 + assert turn2_tools[0].id == "functions.get_news:0" + + +def _split_tool_output_to_deltas( + content: str, tool_strs: list[tuple[str, str]] +) -> list[str]: + """Build a list of string deltas with special tokens as separate chunks. + + Args: + content: text before tool section + tool_strs: list of (tool_id, args_json) + """ + deltas = [content, SECTION_BEGIN] + for tool_id, args_json in tool_strs: + deltas.extend( + [ + TOOL_BEGIN, + f"{tool_id} ", + ARG_BEGIN, + f"{args_json} ", + TOOL_END, + ] + ) + deltas.append(SECTION_END) + return deltas + + +class TestStreamingHappyPath: + def test_single_tool_call(self, parser): + """Verify DeltaToolCall output: name, id, arguments for one tool.""" + deltas = _split_tool_output_to_deltas( + "I'll help. ", + [("functions.get_weather:0", '{"city": "Beijing"}')], + ) + rec = run_tool_extraction_streaming(parser, deltas) + + assert len(rec.tool_calls) == 1 + tc = rec.tool_calls[0] + assert tc.function.name == "get_weather" + assert tc.id == "functions.get_weather:0" + assert json.loads(tc.function.arguments) == {"city": "Beijing"} + + def test_multiple_tool_calls(self, parser): + """Two tool calls emitted with correct indices, names, arguments.""" + deltas = _split_tool_output_to_deltas( + "Compare weather. ", + [ + ("functions.get_weather:0", '{"city": "Tokyo"}'), + ("functions.get_weather:1", '{"city": "NYC"}'), + ], + ) + rec = run_tool_extraction_streaming(parser, deltas) + + assert len(rec.tool_calls) == 2 + assert rec.tool_calls[0].function.name == "get_weather" + assert rec.tool_calls[0].id == "functions.get_weather:0" + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"} + + assert rec.tool_calls[1].function.name == "get_weather" + assert rec.tool_calls[1].id == "functions.get_weather:1" + assert json.loads(rec.tool_calls[1].function.arguments) == {"city": "NYC"} + + def test_content_before_tools(self, parser): + """Content before section is streamed; markers/args don't leak.""" + deltas = _split_tool_output_to_deltas( + "I'll check the weather. ", + [("functions.get_weather:0", '{"city": "Tokyo"}')], + ) + rec = run_tool_extraction_streaming(parser, deltas) + + assert "check the weather" in rec.other_content + # No markers or tool content leaked + for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]: + assert marker not in rec.other_content + assert "get_weather" not in rec.other_content + assert "Tokyo" not in rec.other_content + + def test_no_tool_calls(self, parser): + """Plain text streaming returns content only.""" + deltas = ["This is just ", "regular text ", "without tools."] + rec = run_tool_extraction_streaming(parser, deltas) + + assert rec.other_content == "This is just regular text without tools." + assert rec.tool_calls == [] + + def test_incremental_arguments(self, parser): + """Arguments split across small chunks accumulate correctly.""" + deltas = [ + "Help. ", + SECTION_BEGIN, + TOOL_BEGIN, + "functions.get_weather:0 ", + ARG_BEGIN, + '{"ci', + 'ty": "Be', + 'ijing"}', + " ", + TOOL_END, + SECTION_END, + ] + rec = run_tool_extraction_streaming(parser, deltas) + + assert len(rec.tool_calls) == 1 + assert rec.tool_calls[0].function.name == "get_weather" + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"} + + @pytest.mark.parametrize( + "model_output", + [ + pytest.param( + "Single. " + + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')), + id="single_tool", + ), + pytest.param( + "Multi. " + + _wrap( + _tool("functions.get_weather:0", '{"city": "Tokyo"}'), + _tool("functions.get_news:1", '{"topic": "tech"}'), + ), + id="parallel_tools", + ), + pytest.param( + "No prefix id. " + _wrap(_tool("get_weather:0", '{"city": "NYC"}')), + id="no_functions_prefix", + ), + ], + ) + def test_streaming_matches_nonstreaming(self, parser, model_output): + """Streaming reconstruction matches non-streaming extraction.""" + content_non, tools_non = run_tool_extraction( + parser, model_output, streaming=False + ) + content_stream, tools_stream = run_tool_extraction( + parser, model_output, streaming=True ) + assert len(tools_non) == len(tools_stream) + for tc_non, tc_stream in zip(tools_non, tools_stream): + assert tc_non.function.name == tc_stream.function.name + assert json.loads(tc_non.function.arguments) == json.loads( + tc_stream.function.arguments + ) -def run_streaming_sequence(parser, deltas): - """Helper to simulate a streaming sequence and return results.""" - previous_text = "" - previous_token_ids: list[int] = [] - results = [] - for delta_text, delta_token_ids in deltas: - current_text = previous_text + delta_text - current_token_ids = previous_token_ids + delta_token_ids - - result = parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=delta_text, - previous_token_ids=previous_token_ids, - current_token_ids=current_token_ids, - delta_token_ids=delta_token_ids, - request=None, +class TestStreamingEdgeCases: + def test_marker_suppression(self, parser): + """No special-token markers appear in reconstructed content.""" + deltas = _split_tool_output_to_deltas( + "I'll check. ", + [("functions.get_weather:0", '{"city": "Tokyo"}')], ) - results.append(result) + rec = run_tool_extraction_streaming(parser, deltas) - previous_text = current_text - previous_token_ids = current_token_ids + forbidden = [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN] + for marker in forbidden: + assert marker not in rec.other_content, ( + f"Marker leaked: {marker!r} in {rec.other_content!r}" + ) - return results + def test_noise_between_markers_suppressed(self, parser): + """Text between section_begin and tool_call_begin doesn't leak.""" + deltas = [ + "Reasoning. ", + SECTION_BEGIN, + " spurious noise ", + TOOL_BEGIN, + "functions.test:0 ", + ARG_BEGIN, + '{"k": "v"} ', + TOOL_END, + SECTION_END, + ] + rec = run_tool_extraction_streaming(parser, deltas) + assert "spurious" not in rec.other_content + assert "noise" not in rec.other_content -def test_extract_tool_calls_no_tools(kimi_k2_tool_parser): - model_output = "This is a test" - extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - assert not extracted_tool_calls.tools_called - assert extracted_tool_calls.tool_calls == [] - assert extracted_tool_calls.content == model_output + def test_empty_tool_section(self, parser): + """Empty section (begin immediately followed by end) doesn't crash.""" + deltas = ["Reasoning. ", SECTION_BEGIN, SECTION_END] + rec = run_tool_extraction_streaming(parser, deltas) + assert rec.tool_calls == [] -@pytest.mark.parametrize( - ids=[ - "tool_call_with_content_before", - "multi_tool_call_with_content_before", - "concatenated_tool_calls_bug_fix", - "three_concatenated_tool_calls", - "mixed_spacing_tool_calls", - "angle_brackets_in_json", - "newlines_in_json", - ], - argnames=["model_output", "expected_tool_calls", "expected_content"], - argvalues=[ - ( - """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|> -functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_calls_section_end|>""", + def test_three_different_tools(self, parser): + """Three tool calls with different functions stream correctly.""" + deltas = _split_tool_output_to_deltas( + "Multiple tasks. ", [ - ToolCall( - id="functions.get_weather:0", - function=FunctionCall( - name="get_weather", - arguments=json.dumps( - { - "city": "Beijing", - }, - ), - ), - type="function", - ) + ("functions.get_weather:0", '{"city": "NYC"}'), + ("functions.get_news:1", '{"topic": "tech"}'), + ("functions.send_email:2", '{"to": "a@b.com"}'), ], - "I'll help you check the weather. ", - ), - ( - """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|> -functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_call_begin|> -functions.get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>""", - [ - ToolCall( - id="functions.get_weather:0", - function=FunctionCall( - name="get_weather", - arguments=json.dumps( - { - "city": "Beijing", - }, - ), - ), - type="function", - ), - ToolCall( - id="functions.get_weather:1", - function=FunctionCall( - name="get_weather", - arguments=json.dumps( - { - "city": "Shanghai", - }, - ), - ), - type="function", - ), - ], - "I'll help you check the weather. ", - ), - ( - """I'll get the weather and news for LA today. First, let me get the weather using Los Angeles coordinates, and then get the latest news. <|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"latitude": 34.0522, "longitude": -118.2437}<|tool_call_end|><|tool_call_begin|>functions.get_news:1<|tool_call_argument_begin|>{"content": "Los Angeles today"}<|tool_call_end|><|tool_calls_section_end|>""", - [ - ToolCall( - id="functions.get_weather:0", - function=FunctionCall( - name="get_weather", - arguments=json.dumps( - {"latitude": 34.0522, "longitude": -118.2437} - ), - ), - type="function", - ), - ToolCall( - id="functions.get_news:1", - function=FunctionCall( - name="get_news", - arguments=json.dumps({"content": "Los Angeles today"}), - ), - type="function", - ), - ], - "I'll get the weather and news for LA today. First, let me get the weather using Los Angeles coordinates, and then get the latest news. ", - ), - ( - """I'll help you with multiple tasks. <|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"city": "New York"}<|tool_call_end|><|tool_call_begin|>functions.get_news:1<|tool_call_argument_begin|>{"topic": "technology"}<|tool_call_end|><|tool_call_begin|>functions.send_email:2<|tool_call_argument_begin|>{"to": "user@example.com", "subject": "Daily Update"}<|tool_call_end|><|tool_calls_section_end|>""", - [ - ToolCall( - id="functions.get_weather:0", - function=FunctionCall( - name="get_weather", - arguments=json.dumps({"city": "New York"}), - ), - type="function", - ), - ToolCall( - id="functions.get_news:1", - function=FunctionCall( - name="get_news", - arguments=json.dumps({"topic": "technology"}), - ), - type="function", - ), - ToolCall( - id="functions.send_email:2", - function=FunctionCall( - name="send_email", - arguments=json.dumps( - {"to": "user@example.com", "subject": "Daily Update"} - ), - ), - type="function", - ), - ], - "I'll help you with multiple tasks. ", - ), - ( - """Mixed spacing test. <|tool_calls_section_begin|> <|tool_call_begin|> functions.test:0 <|tool_call_argument_begin|> {} <|tool_call_end|><|tool_call_begin|>functions.test2:1<|tool_call_argument_begin|>{}<|tool_call_end|> <|tool_calls_section_end|>""", - [ - ToolCall( - id="functions.test:0", - function=FunctionCall( - name="test", - arguments=json.dumps({}), - ), - type="function", - ), - ToolCall( - id="functions.test2:1", - function=FunctionCall( - name="test2", - arguments=json.dumps({}), - ), - type="function", - ), - ], - "Mixed spacing test. ", - ), - ( - """I need to process HTML content. <|tool_calls_section_begin|><|tool_call_begin|>functions.process_html:0<|tool_call_argument_begin|>{"html": "
content
", "text": "normal text"}<|tool_call_end|><|tool_calls_section_end|>""", - [ - ToolCall( - id="functions.process_html:0", - function=FunctionCall( - name="process_html", - arguments=json.dumps( - {"html": "
content
", "text": "normal text"} - ), - ), - type="function", - ) - ], - "I need to process HTML content. ", - ), - ( - """I need to process formatted JSON. <|tool_calls_section_begin|><|tool_call_begin|>functions.process_data:0<|tool_call_argument_begin|>{ - "name": "test", - "value": 123, - "nested": { - "key": "value" - } -}<|tool_call_end|><|tool_calls_section_end|>""", - [ - ToolCall( - id="functions.process_data:0", - function=FunctionCall( - name="process_data", - arguments=json.dumps( - {"name": "test", "value": 123, "nested": {"key": "value"}}, - indent=2, - ), - ), - type="function", - ) - ], - "I need to process formatted JSON. ", - ), - ], -) -def test_extract_tool_calls( - kimi_k2_tool_parser, model_output, expected_tool_calls, expected_content -): - extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - assert extracted_tool_calls.tools_called - - assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls) - - assert extracted_tool_calls.content == expected_content - - -def test_extract_tool_calls_invalid_json(kimi_k2_tool_parser): - """we'll return every funcall result""" - model_output = """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|> -functions.invalid_get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing" <|tool_call_end|> <|tool_call_begin|> -functions.valid_get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>""" - - extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - - assert extracted_tool_calls.tools_called - # Should extract only the valid JSON tool calls - assert len(extracted_tool_calls.tool_calls) == 2 - assert extracted_tool_calls.tool_calls[0].function.name == "invalid_get_weather" - assert extracted_tool_calls.tool_calls[1].function.name == "valid_get_weather" - - -def test_extract_tool_calls_invalid_funcall(kimi_k2_tool_parser): - """we'll return every funcall result""" - model_output = """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|> -functions.invalid_get_weather.0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_call_begin|> -functions.valid_get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>""" - - extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - - assert extracted_tool_calls.tools_called - # Should extract only the valid JSON tool calls - assert len(extracted_tool_calls.tool_calls) == 1 - assert extracted_tool_calls.tool_calls[0].function.name == "valid_get_weather" - - -def test_streaming_basic_functionality(kimi_k2_tool_parser): - """Test basic streaming functionality.""" - # Reset streaming state - kimi_k2_tool_parser.current_tool_name_sent = False - kimi_k2_tool_parser.prev_tool_call_arr = [] - kimi_k2_tool_parser.current_tool_id = -1 - kimi_k2_tool_parser.streamed_args_for_tool = [] - - # Test with a simple tool call - current_text = """ check the weather. <|tool_calls_section_begin|> <|tool_call_begin|> -functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_calls_section_end|>""" - - # First call should handle the initial setup - result = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="I'll help you", - current_text=current_text, - delta_text="<|tool_calls_section_end|>", - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - # The result might be None or contain tool call information - # This depends on the internal state management - if result is not None and hasattr(result, "tool_calls") and result.tool_calls: - assert len(result.tool_calls) >= 0 - - -def test_streaming_no_tool_calls(kimi_k2_tool_parser): - """Test streaming when there are no tool calls.""" - current_text = "This is just regular text without any tool calls." - - result = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="This is just regular text", - current_text=current_text, - delta_text=" without any tool calls.", - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - # Should return the delta text as content - assert result is not None - assert hasattr(result, "content") - assert result.content == " without any tool calls." - - -def test_token_leak_between_section_and_tool_begin(kimi_k2_tool_parser): - """ - Test that text between <|tool_calls_section_begin|> and <|tool_call_begin|> - is suppressed and does not leak into reasoning_delta. - This is the main vulnerability being fixed. - """ - kimi_k2_tool_parser.reset_streaming_state() - - # Get token IDs for the markers - section_begin_token_id = kimi_k2_tool_parser.vocab.get( - "<|tool_calls_section_begin|>" - ) - tool_call_begin_token_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>") - - # Simulate streaming sequence: - deltas = [ - ("I'll help you with that. ", [1, 2, 3]), - ("<|tool_calls_section_begin|>", [section_begin_token_id]), - (" spurious text ", [4, 5]), - ("<|tool_call_begin|>", [tool_call_begin_token_id]), - ] - - results = run_streaming_sequence(kimi_k2_tool_parser, deltas) - - # Delta 1: "I'll help you with that. " - assert results[0] is not None - assert results[0].content == "I'll help you with that. " - - # Delta 2: "<|tool_calls_section_begin|>" - # Section marker should be stripped and suppressed - assert results[1] is None or ( - results[1].content is None or results[1].content == "" - ) - - # Delta 3: " spurious text or tokens " (THE LEAK SCENARIO) - # CRITICAL: This text should be suppressed, NOT returned as reasoning_delta - assert results[2] is None or ( - results[2].content is None or results[2].content == "" - ) - - # Delta 4: "<|tool_call_begin|>..." - # Now we're in tool call mode, result depends on internal state - # The key is that the spurious text from Delta 3 was not leaked - - -def test_split_markers_across_deltas(kimi_k2_tool_parser): - """ - Test that markers split across delta chunks are correctly detected - via the rolling buffer mechanism. - """ - kimi_k2_tool_parser.reset_streaming_state() - - section_begin_token_id = kimi_k2_tool_parser.vocab.get( - "<|tool_calls_section_begin|>" - ) - - # Delta 1: partial token, Delta 2: complete marker - deltas = [ - ("<|tool_calls_sec", [3]), - ("tion_begin|> ", [section_begin_token_id, 4]), - ] - - _results = run_streaming_sequence(kimi_k2_tool_parser, deltas) - - # Now the complete marker should be detected via buffer - assert kimi_k2_tool_parser.in_tool_section is True - - -def test_marker_variants(kimi_k2_tool_parser): - """Test that both singular and plural marker variants are recognized.""" - kimi_k2_tool_parser.reset_streaming_state() - - # Test singular variant: <|tool_call_section_begin|> (note: singular "call") - singular_token_id = kimi_k2_tool_parser.vocab.get("<|tool_call_section_begin|>") - - if singular_token_id is not None: # Only test if tokenizer supports it - _result = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="Reasoning ", - current_text="Reasoning <|tool_call_section_begin|>", - delta_text="<|tool_call_section_begin|>", - previous_token_ids=[1, 2], - current_token_ids=[1, 2, singular_token_id], - delta_token_ids=[singular_token_id], - request=None, ) - # Should enter tool section mode with singular variant too - assert kimi_k2_tool_parser.in_tool_section is True - - -def test_reentry_to_reasoning_after_tool_section(kimi_k2_tool_parser): - """ - Test that after exiting a tool section with <|tool_calls_section_end|>, - subsequent text is correctly returned as reasoning content. - """ - kimi_k2_tool_parser.reset_streaming_state() - - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>") - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>") - - deltas = [ - ("<|tool_calls_section_begin|>", [section_begin_id]), - ("<|tool_calls_section_end|>", [section_end_id]), - (" More reasoning", [10, 11]), - ] - - results = run_streaming_sequence(kimi_k2_tool_parser, deltas) - - assert kimi_k2_tool_parser.in_tool_section is False - assert results[2] is not None - assert results[2].content == " More reasoning" - - -def test_empty_tool_section(kimi_k2_tool_parser): - """Test an empty tool section (begin immediately followed by end).""" - kimi_k2_tool_parser.reset_streaming_state() - - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>") - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>") - - # Section begin - _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="Reasoning ", - current_text="Reasoning <|tool_calls_section_begin|>", - delta_text="<|tool_calls_section_begin|>", - previous_token_ids=[1], - current_token_ids=[1, section_begin_id], - delta_token_ids=[section_begin_id], - request=None, - ) - - # Immediate section end - _result2 = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="Reasoning <|tool_calls_section_begin|>", - current_text="Reasoning <|tool_calls_section_begin|><|tool_calls_section_end|>", - delta_text="<|tool_calls_section_end|>", - previous_token_ids=[1, section_begin_id], - current_token_ids=[1, section_begin_id, section_end_id], - delta_token_ids=[section_end_id], - request=None, - ) - # Should exit cleanly without errors - assert kimi_k2_tool_parser.in_tool_section is False - - -def test_malformed_tool_section_recovery(kimi_k2_tool_parser): - """ - Test that the parser recovers from a malformed tool section - that never closes properly. - """ - kimi_k2_tool_parser.reset_streaming_state() - - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>") - - # Enter tool section - _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text="<|tool_calls_section_begin|>", - delta_text="<|tool_calls_section_begin|>", - previous_token_ids=[], - current_token_ids=[section_begin_id], - delta_token_ids=[section_begin_id], - request=None, - ) - assert kimi_k2_tool_parser.in_tool_section is True - - # Simulate a lot of text without proper tool calls or section end - # This should trigger the error recovery mechanism - large_text = "x" * 10000 # Exceeds max_section_chars - - result2 = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="<|tool_calls_section_begin|>", - current_text="<|tool_calls_section_begin|>" + large_text, - delta_text=large_text, - previous_token_ids=[section_begin_id], - current_token_ids=[section_begin_id] + list(range(100, 100 + len(large_text))), - delta_token_ids=list(range(100, 100 + len(large_text))), - request=None, - ) - - # Parser should have force-exited the tool section - assert kimi_k2_tool_parser.in_tool_section is False - # And returned the content as reasoning - assert result2 is not None - assert result2.content == large_text - - -def test_state_reset(kimi_k2_tool_parser): - """Test that reset_streaming_state() properly clears all state.""" - # Put parser in a complex state - kimi_k2_tool_parser.in_tool_section = True - kimi_k2_tool_parser.token_buffer = "some buffer" - kimi_k2_tool_parser.current_tool_id = 5 - kimi_k2_tool_parser.prev_tool_call_arr = [{"id": "test"}] - kimi_k2_tool_parser.section_char_count = 1000 - - # Reset - kimi_k2_tool_parser.reset_streaming_state() - - # Verify all state is cleared - assert kimi_k2_tool_parser.in_tool_section is False - assert kimi_k2_tool_parser.token_buffer == "" - assert kimi_k2_tool_parser.current_tool_id == -1 - assert kimi_k2_tool_parser.prev_tool_call_arr == [] - assert kimi_k2_tool_parser.section_char_count == 0 - assert kimi_k2_tool_parser.current_tool_name_sent is False - assert kimi_k2_tool_parser.streamed_args_for_tool == [] - - -def test_section_begin_noise_tool_begin_same_chunk(kimi_k2_tool_parser): - """ - Test that begin→noise→tool_begin within the SAME chunk suppresses - the noise text correctly (not just across chunks). - """ - kimi_k2_tool_parser.reset_streaming_state() - - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>") - tool_call_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>") - - # Single delta containing: section_begin + spurious text + tool_call_begin - combined_text = "<|tool_calls_section_begin|> noise text <|tool_call_begin|>" - - result = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="Reasoning ", - current_text="Reasoning " + combined_text, - delta_text=combined_text, - previous_token_ids=[1, 2], - current_token_ids=[1, 2, section_begin_id, 3, 4, tool_call_begin_id], - delta_token_ids=[section_begin_id, 3, 4, tool_call_begin_id], - request=None, - ) - - # The noise text should NOT leak into content - # Result should either be None/empty or start tool call parsing - if result is not None and result.content is not None: - # If content is returned, it should not contain the noise - assert "noise text" not in result.content - assert result.content == "" or result.content.strip() == "" - - -def test_stream_ends_without_section_end_marker(kimi_k2_tool_parser): - """ - Test that if the stream ends (EOF) without a proper section end marker, - the parser doesn't leak text, doesn't crash, and resets state cleanly. - """ - kimi_k2_tool_parser.reset_streaming_state() - - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>") - - # Enter tool section - _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text="<|tool_calls_section_begin|>", - delta_text="<|tool_calls_section_begin|>", - previous_token_ids=[], - current_token_ids=[section_begin_id], - delta_token_ids=[section_begin_id], - request=None, - ) - assert kimi_k2_tool_parser.in_tool_section is True - - # Some content in tool section - result2 = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="<|tool_calls_section_begin|>", - current_text="<|tool_calls_section_begin|> partial content", - delta_text=" partial content", - previous_token_ids=[section_begin_id], - current_token_ids=[section_begin_id, 10, 11], - delta_token_ids=[10, 11], - request=None, - ) - # Content should be suppressed - assert result2.content == "" or result2.content is None - - # Stream ends (EOF) - no more deltas, no section_end marker - # Simulate this by manually checking state and resetting - # (In real usage, the request handler would call reset_streaming_state) - assert kimi_k2_tool_parser.in_tool_section is True # Still in section - - # Reset state (as would happen between requests) - kimi_k2_tool_parser.reset_streaming_state() - - # Verify clean slate - assert kimi_k2_tool_parser.in_tool_section is False - assert kimi_k2_tool_parser.token_buffer == "" - - # Next request should work normally - result3 = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text="New reasoning", - delta_text="New reasoning", - previous_token_ids=[], - current_token_ids=[20, 21], - delta_token_ids=[20, 21], - request=None, - ) - assert result3 is not None - assert result3.content == "New reasoning" - - -def test_same_chunk_begin_and_end_markers(kimi_k2_tool_parser): - """ - CRITICAL TEST: Verify that when both section_begin and section_end - markers appear in the SAME chunk, the parser correctly: - 1. Enters the tool section - 2. Immediately exits the tool section - 3. Does NOT get stuck in in_tool_section=True state - - This tests the bug fix where elif was changed to if to handle - both state transitions in a single delta. - """ - kimi_k2_tool_parser.reset_streaming_state() - - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>") - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>") - - # Single chunk with both markers (e.g., empty tool section) - combined_delta = "<|tool_calls_section_begin|><|tool_calls_section_end|>" - - result = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="Some reasoning ", - current_text="Some reasoning " + combined_delta, - delta_text=combined_delta, - previous_token_ids=[1, 2], - current_token_ids=[1, 2, section_begin_id, section_end_id], - delta_token_ids=[section_begin_id, section_end_id], - request=None, - ) - - # CRITICAL: Parser should NOT be stuck in tool section - assert kimi_k2_tool_parser.in_tool_section is False, ( - "Parser stuck in tool section after processing both begin/end in same chunk. " - "This indicates the elif bug was not fixed." - ) - - # Result should be empty or contain only stripped content - assert result is not None - assert result.content == "" or result.content is None - - # Verify subsequent content streams correctly (not suppressed) - result2 = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="Some reasoning " + combined_delta, - current_text="Some reasoning " + combined_delta + " More reasoning", - delta_text=" More reasoning", - previous_token_ids=[1, 2, section_begin_id, section_end_id], - current_token_ids=[1, 2, section_begin_id, section_end_id, 10, 11], - delta_token_ids=[10, 11], - request=None, - ) - - # This content should NOT be suppressed (we're out of tool section) - assert result2 is not None - assert result2.content == " More reasoning" - - -def test_same_chunk_begin_content_end_markers(kimi_k2_tool_parser): - """ - Test the same-chunk scenario with actual content between markers. - Example: <|tool_calls_section_begin|> text <|tool_calls_section_end|> - all arriving in one delta. The key is that the state machine correctly - transitions in and out within the same chunk. - """ - kimi_k2_tool_parser.reset_streaming_state() - - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>") - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>") - - # Chunk with begin, some whitespace/noise, and end all together - # This simulates a tool section that opens and closes in the same chunk - combined_delta = "<|tool_calls_section_begin|> <|tool_calls_section_end|>" - - _result = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="Reasoning ", - current_text="Reasoning " + combined_delta, - delta_text=combined_delta, - previous_token_ids=[1], - current_token_ids=[1, section_begin_id, 100, section_end_id], - delta_token_ids=[section_begin_id, 100, section_end_id], - request=None, - ) - - # Parser should exit cleanly (not stuck in tool section) - assert kimi_k2_tool_parser.in_tool_section is False - - # Verify the fix: next content should stream normally, not be suppressed - result2 = kimi_k2_tool_parser.extract_tool_calls_streaming( - previous_text="Reasoning " + combined_delta, - current_text="Reasoning " + combined_delta + " Done", - delta_text=" Done", - previous_token_ids=[1, section_begin_id, 100, section_end_id], - current_token_ids=[1, section_begin_id, 100, section_end_id, 200], - delta_token_ids=[200], - request=None, - ) - - # Content after section should be returned (not suppressed) - assert result2 is not None - assert result2.content == " Done" - - -def test_tool_call_end_and_section_end_same_chunk(kimi_k2_tool_parser): - """ - CRITICAL TEST (P1): Verify that when both <|tool_call_end|> and - <|tool_calls_section_end|> appear in the SAME chunk, the parser: - 1. Processes the tool_call_end first (emits final arguments) - 2. THEN exits the section - 3. Does NOT drop the final tool call update - 4. Does NOT leak special tokens into reasoning - - This tests the deferred section exit fix. - """ - kimi_k2_tool_parser.reset_streaming_state() - - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>") - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>") - tool_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>") - tool_end_id = kimi_k2_tool_parser.vocab.get("<|tool_call_end|>") - - # Simulate a streaming sequence for a SHORT tool call (all in one chunk): - combined = ( - '<|tool_call_begin|>get_weather:0 <|tool_call_argument_begin|> {"city": "Paris"} ' - "<|tool_call_end|><|tool_calls_section_end|>" - ) - - deltas = [ - ("Let me help. ", [1, 2]), - ("<|tool_calls_section_begin|>", [section_begin_id]), - (combined, [tool_begin_id, 10, 11, 12, tool_end_id, section_end_id]), - (" Done", [20]), - ] - - results = run_streaming_sequence(kimi_k2_tool_parser, deltas) - - # CRITICAL: Parser should have exited section AFTER processing tool - assert kimi_k2_tool_parser.in_tool_section is False - - # Tool call should have been emitted (not dropped) - if results[2] is not None and results[2].content is not None: - # Verify no special tokens leaked into content - assert "<|tool_call_end|>" not in results[2].content - assert "<|tool_calls_section_end|>" not in results[2].content - - # Content after tool section should stream normally - assert results[3] is not None - assert results[3].content == " Done" - - -def test_streaming_tool_call_markers_not_leaked(kimi_k2_tool_parser): - """ - CRITICAL TEST: Verify that tool call markers (<|tool_call_begin|>, - <|tool_call_end|>, <|tool_call_argument_begin|>) are NOT leaked - into the content field during streaming. - - This reproduces the AWS Bedrock bug where tool call markers appeared - in the 'text' field of responses. - """ - kimi_k2_tool_parser.reset_streaming_state() - - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>") - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>") - tool_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>") - tool_end_id = kimi_k2_tool_parser.vocab.get("<|tool_call_end|>") - - # List of markers that should NEVER appear in content - forbidden_markers = [ - "<|tool_call_begin|>", - "<|tool_call_end|>", - "<|tool_call_argument_begin|>", - "<|tool_calls_section_begin|>", - "<|tool_calls_section_end|>", - ] - - all_content = [] - - # Steps: reasoning, section begin, tool call, section end, more reasoning - tool_chunk = ( - "<|tool_call_begin|> functions.get_weather:0 " - '<|tool_call_argument_begin|> {"city": "Tokyo"} <|tool_call_end|>' - ) - deltas = [ - ("I'll check the weather. ", [1, 2, 3]), - ("<|tool_calls_section_begin|>", [section_begin_id]), - (tool_chunk, [tool_begin_id, 10, 11, tool_end_id]), - ("<|tool_calls_section_end|>", [section_end_id]), - (" Here's the result.", [20, 21]), - ] - - results = run_streaming_sequence(kimi_k2_tool_parser, deltas) - - for res in results: - if res and res.content: - all_content.append(res.content) - - # CRITICAL ASSERTIONS: No forbidden markers in any content - full_content = "".join(all_content) - for marker in forbidden_markers: - assert marker not in full_content, ( - f"MARKER LEAK DETECTED: '{marker}' found in content. " - f"Full content: {repr(full_content)}" + rec = run_tool_extraction_streaming(parser, deltas) + + assert len(rec.tool_calls) == 3 + names = [tc.function.name for tc in rec.tool_calls] + assert names == ["get_weather", "get_news", "send_email"] + ids = [tc.id for tc in rec.tool_calls] + assert len(set(ids)) == 3 # unique ids + + def test_truncated_tool_call_no_end_marker(self, parser): + """Stream ending mid-tool-call (max_tokens) doesn't crash.""" + deltas = [ + "I'll check. ", + SECTION_BEGIN, + TOOL_BEGIN, + "functions.get_weather:0 ", + ARG_BEGIN, + '{"city": "Bei', + # Stream ends here — no TOOL_END, no SECTION_END + ] + rec = run_tool_extraction_streaming(parser, deltas) + + # Should not crash; tool name and partial args extracted + assert len(rec.tool_calls) == 1 + assert rec.tool_calls[0].function.name == "get_weather" + assert rec.tool_calls[0].id == "functions.get_weather:0" + assert rec.tool_calls[0].function.arguments == '{"city": "Bei' + # No markers leaked into content + for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]: + assert marker not in rec.other_content + + def test_content_after_tool_section(self, parser): + """Trailing text after section_end doesn't crash or leak markers.""" + deltas = [ + "Before. ", + SECTION_BEGIN, + TOOL_BEGIN, + "functions.get_weather:0 ", + ARG_BEGIN, + '{"city": "Tokyo"} ', + TOOL_END, + SECTION_END, + " After tools.", + ] + rec = run_tool_extraction_streaming(parser, deltas) + + # Tool call extracted correctly + assert len(rec.tool_calls) == 1 + assert rec.tool_calls[0].function.name == "get_weather" + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"} + # Trailing content after tool section is dropped + assert "After tools." not in rec.other_content + # No markers leaked into content + for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]: + assert marker not in rec.other_content + + +class TestAdjustRequest: + def test_sets_skip_special_tokens_false(self, parser): + request = MagicMock(spec=ChatCompletionRequest) + request.tools = [{"type": "function", "function": {"name": "test"}}] + request.tool_choice = "auto" + request.skip_special_tokens = True + + result = parser.adjust_request(request) + assert result.skip_special_tokens is False + + def test_no_change_when_tool_choice_none(self, parser): + request = MagicMock(spec=ChatCompletionRequest) + request.tools = [{"type": "function", "function": {"name": "test"}}] + request.tool_choice = "none" + request.skip_special_tokens = True + + result = parser.adjust_request(request) + assert result.skip_special_tokens is True + + def test_no_change_when_no_tools(self, parser): + request = MagicMock(spec=ChatCompletionRequest) + request.tools = None + request.tool_choice = "auto" + request.skip_special_tokens = True + + result = parser.adjust_request(request) + assert result.skip_special_tokens is True + + +def _chunk_tokenized_deltas(tokenizer, text: str, stream_interval: int) -> list[str]: + """Encode text, group tokens into chunks of stream_interval, decode each.""" + token_ids = tokenizer.encode(text, add_special_tokens=False) + deltas = [] + prev = "" + for i in range(0, len(token_ids), stream_interval): + decoded = tokenizer.decode( + token_ids[: i + stream_interval], skip_special_tokens=False + ) + deltas.append(decoded[len(prev) :]) + prev = decoded + return deltas + + +class TestStreamingIntervals: + """Test streaming at various token-chunk sizes to catch boundary bugs.""" + + @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) + def test_single_tool_call_at_interval(self, kimi_k2_tokenizer, stream_interval): + text = "Help. " + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')) + deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval) + parser = KimiK2ToolParser(kimi_k2_tokenizer) + rec = run_tool_extraction_streaming( + parser, deltas, assert_one_tool_per_delta=False ) - # Also check that tool call content (function name, arguments) is not leaked - assert "get_weather" not in full_content, ( - f"TOOL CALL CONTENT LEAKED: 'get_weather' found in content. " - f"Full content: {repr(full_content)}" - ) - assert "Tokyo" not in full_content, ( - f"TOOL CALL CONTENT LEAKED: 'Tokyo' found in content. " - f"Full content: {repr(full_content)}" - ) + assert len(rec.tool_calls) == 1 + assert rec.tool_calls[0].function.name == "get_weather" + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"} - # Verify that legitimate content was preserved - assert "I'll check the weather." in full_content or len(all_content) > 0 - - -def test_native_id_extracted_and_placed_on_tool_call(kimi_k2_tool_parser): - """Regression: parser extracts native ID onto ToolCall (PR #32768).""" - model_output = ( - "Checking weather. " - "<|tool_calls_section_begin|>" - "<|tool_call_begin|>functions.get_weather:0" - '<|tool_call_argument_begin|>{"city": "Tokyo"}' - "<|tool_call_end|>" - "<|tool_calls_section_end|>" - ) - - result = kimi_k2_tool_parser.extract_tool_calls(model_output, request=None) - assert result.tools_called - assert len(result.tool_calls) == 1 - - tc = result.tool_calls[0] - # Native ID from model output must be used as the tool call ID - assert tc.id == "functions.get_weather:0" - assert tc.function.name == "get_weather" - assert json.loads(tc.function.arguments) == {"city": "Tokyo"} - - -def test_multi_turn_native_id_continuity(kimi_k2_tool_parser, kimi_k2_tokenizer): - """Regression: native IDs from turn 1 preserved across turns (PR #32768).""" - turn1_output = ( - "Let me check. " - "<|tool_calls_section_begin|>" - "<|tool_call_begin|>functions.get_weather:0" - '<|tool_call_argument_begin|>{"city": "Beijing"}' - "<|tool_call_end|>" - "<|tool_calls_section_end|>" - ) - - turn1_result = kimi_k2_tool_parser.extract_tool_calls(turn1_output, request=None) - assert turn1_result.tools_called - assert turn1_result.tool_calls[0].id == "functions.get_weather:0" - - # Fresh parser for turn 2 - turn2_parser = KimiK2ToolParser(kimi_k2_tokenizer) - turn2_output = ( - "Now let me get news. " - "<|tool_calls_section_begin|>" - "<|tool_call_begin|>functions.get_news:0" - '<|tool_call_argument_begin|>{"topic": "weather in Beijing"}' - "<|tool_call_end|>" - "<|tool_calls_section_end|>" - ) - - turn2_result = turn2_parser.extract_tool_calls(turn2_output, request=None) - assert turn2_result.tools_called - assert turn2_result.tool_calls[0].id == "functions.get_news:0" - - -def test_streaming_multiple_tool_calls_not_leaked(kimi_k2_tool_parser): - """ - Test that MULTIPLE tool calls in streaming mode do not leak into content. - This reproduces the AWS Bedrock scenario: "Compare weather in Tokyo and NYC". - """ - kimi_k2_tool_parser.reset_streaming_state() - - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>") - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>") - tool_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>") - tool_end_id = kimi_k2_tool_parser.vocab.get("<|tool_call_end|>") - - all_content = [] - - tool1 = '<|tool_call_begin|> get_weather:0 <|tool_call_argument_begin|> {"city": "Tokyo"} <|tool_call_end|>' - tool2 = ' <|tool_call_begin|> get_weather:1 <|tool_call_argument_begin|> {"city": "New York"} <|tool_call_end|>' - - deltas = [ - ("I'll compare the weather. ", [1, 2, 3]), - ("<|tool_calls_section_begin|>", [section_begin_id]), - (tool1, [tool_begin_id, 10, tool_end_id]), - (tool2, [tool_begin_id, 20, tool_end_id]), - ("<|tool_calls_section_end|>", [section_end_id]), - (" Here's the comparison.", [30]), - ] - - results = run_streaming_sequence(kimi_k2_tool_parser, deltas) - - for res in results: - if res and res.content: - all_content.append(res.content) - - # Assertions - full_content = "".join(all_content) - - # Check no markers leaked - forbidden = ["<|tool_call", "<|tool_calls_section"] - for marker in forbidden: - assert marker not in full_content, ( - f"MARKER LEAKED: {marker} in {repr(full_content)}" + @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) + def test_content_then_tool_call_at_interval( + self, kimi_k2_tokenizer, stream_interval + ): + text = "Sure, let me check. " + _wrap( + _tool("functions.get_weather:0", '{"city": "Tokyo"}') + ) + deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval) + parser = KimiK2ToolParser(kimi_k2_tokenizer) + rec = run_tool_extraction_streaming( + parser, deltas, assert_one_tool_per_delta=False ) - # Check no tool call content leaked (both tools) - assert "get_weather" not in full_content, f"TOOL NAME LEAKED: {repr(full_content)}" - assert "Tokyo" not in full_content, f"TOOL ARG LEAKED (Tokyo): {repr(full_content)}" - assert "New York" not in full_content, ( - f"TOOL ARG LEAKED (NYC): {repr(full_content)}" - ) + assert "let me check" in rec.other_content + assert "get_weather" not in rec.other_content + assert len(rec.tool_calls) == 1 + assert rec.tool_calls[0].function.name == "get_weather" + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"} - # Legitimate content preserved - assert "compare" in full_content.lower() or len(all_content) > 0 + @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) + def test_multiple_tool_calls_at_interval(self, kimi_k2_tokenizer, stream_interval): + text = "Compare. " + _wrap( + _tool("functions.search:0", '{"q": "cats"}'), + _tool("functions.search:1", '{"q": "dogs"}'), + ) + deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval) + parser = KimiK2ToolParser(kimi_k2_tokenizer) + rec = run_tool_extraction_streaming( + parser, deltas, assert_one_tool_per_delta=False + ) + + assert len(rec.tool_calls) == 2 + assert rec.tool_calls[0].function.name == "search" + assert json.loads(rec.tool_calls[0].function.arguments) == {"q": "cats"} + assert rec.tool_calls[1].function.name == "search" + assert json.loads(rec.tool_calls[1].function.arguments) == {"q": "dogs"} + + @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) + def test_plain_text_at_interval(self, kimi_k2_tokenizer, stream_interval): + text = "This is plain text with no tool calling involved." + deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval) + parser = KimiK2ToolParser(kimi_k2_tokenizer) + rec = run_tool_extraction_streaming( + parser, deltas, assert_one_tool_per_delta=False + ) + + assert rec.other_content == text + assert rec.tool_calls == [] + + def test_content_and_tool_call_in_single_chunk(self, kimi_k2_tokenizer): + """Content + complete tool call in one chunk must both be emitted.""" + text = "Hi! " + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')) + deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval=9999) + parser = KimiK2ToolParser(kimi_k2_tokenizer) + rec = run_tool_extraction_streaming( + parser, deltas, assert_one_tool_per_delta=False + ) + + assert "Hi!" in rec.other_content + assert "get_weather" not in rec.other_content + assert len(rec.tool_calls) == 1 + assert rec.tool_calls[0].function.name == "get_weather" + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"} diff --git a/vllm/tool_parsers/kimi_k2_tool_parser.py b/vllm/tool_parsers/kimi_k2_tool_parser.py index bc995319e51..7ddd8fa7a80 100644 --- a/vllm/tool_parsers/kimi_k2_tool_parser.py +++ b/vllm/tool_parsers/kimi_k2_tool_parser.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# code modified from deepseekv3_tool_parser.py from collections.abc import Sequence @@ -17,12 +16,14 @@ from vllm.entrypoints.openai.engine.protocol import ( FunctionCall, ToolCall, ) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, ) +from vllm.tool_parsers.utils import partial_tag_overlap logger = init_logger(__name__) @@ -30,124 +31,44 @@ logger = init_logger(__name__) class KimiK2ToolParser(ToolParser): def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) - self.current_tool_name_sent: bool = False + + # Streaming state + self._sent_content_idx: int = 0 self.prev_tool_call_arr: list[dict] = [] - self.current_tool_id: int = -1 - self.streamed_args_for_tool: list[ - str - ] = [] # map what has been streamed for each tool so far to a list + self.streamed_args_for_tool: list[str] = [] - # Section-level state management to prevent token leakage - self.in_tool_section: bool = False - self.token_buffer: str = "" - # Buffer size: empirical worst-case for longest marker (~30 chars) * 2 - # + safety margin for unicode + partial overlap. Prevents unbounded growth. - self.buffer_max_size: int = 1024 - self.section_char_count: int = 0 # Track characters processed in tool section - self.max_section_chars: int = 8192 # Force exit if section exceeds this - self._buffer_overflow_logged: bool = False # Log overflow once per session - - # Support both singular and plural variants + # Section marker self.tool_calls_start_token: str = "<|tool_calls_section_begin|>" - self.tool_calls_end_token: str = "<|tool_calls_section_end|>" - self.tool_calls_start_token_variants: list[str] = [ - "<|tool_calls_section_begin|>", - "<|tool_call_section_begin|>", # singular variant - ] - self.tool_calls_end_token_variants: list[str] = [ - "<|tool_calls_section_end|>", - "<|tool_call_section_end|>", # singular variant - ] + # Individual tool call markers self.tool_call_start_token: str = "<|tool_call_begin|>" self.tool_call_end_token: str = "<|tool_call_end|>" + self.tool_call_arg_token: str = "<|tool_call_argument_begin|>" + # Regex for non-streaming extraction self.tool_call_regex = re.compile( - r"<\|tool_call_begin\|>\s*(?P[^<]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P(?:(?!<\|tool_call_begin\|>).)*?)\s*<\|tool_call_end\|>", + r"<\|tool_call_begin\|>\s*(?P[^<]+:\d+)\s*" + r"<\|tool_call_argument_begin\|>\s*" + r"(?P(?:(?!<\|tool_call_begin\|>).)*?)\s*" + r"<\|tool_call_end\|>", re.DOTALL, ) - self.stream_tool_call_portion_regex = re.compile( - r"(?P.+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P.*)" - ) - - self.stream_tool_call_name_regex = re.compile(r"(?P.+:\d+)\s*") - if not self.model_tokenizer: raise ValueError( "The model tokenizer must be passed to the ToolParser " "constructor during construction." ) - self.tool_calls_start_token_id = self.vocab.get(self.tool_calls_start_token) - self.tool_calls_end_token_id = self.vocab.get(self.tool_calls_end_token) - # Get token IDs for all variants - self.tool_calls_start_token_ids: list[int] = [ - tid - for variant in self.tool_calls_start_token_variants - if (tid := self.vocab.get(variant)) is not None - ] - self.tool_calls_end_token_ids: list[int] = [ - tid - for variant in self.tool_calls_end_token_variants - if (tid := self.vocab.get(variant)) is not None - ] - - self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token) - self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token) - - if ( - self.tool_calls_start_token_id is None - or self.tool_calls_end_token_id is None - ): - raise RuntimeError( - "Kimi-K2 Tool parser could not locate tool call start/end " - "tokens in the tokenizer!" - ) - - def _check_and_strip_markers(self, text: str) -> tuple[str, bool, bool]: - """ - Check for section begin/end markers in text and strip them. - Returns: (cleaned_text, found_section_begin, found_section_end) - """ - found_begin = False - found_end = False - cleaned = text - - # Check for section begin markers (any variant) - for variant in self.tool_calls_start_token_variants: - if variant in cleaned: - cleaned = cleaned.replace(variant, "") - found_begin = True - - # Check for section end markers (any variant) - for variant in self.tool_calls_end_token_variants: - if variant in cleaned: - cleaned = cleaned.replace(variant, "") - found_end = True - return cleaned, found_begin, found_end - - def _reset_section_state(self) -> None: - """Reset state when exiting tool section.""" - self.in_tool_section = False - self.token_buffer = "" - self.section_char_count = 0 - - def reset_streaming_state(self) -> None: - """ - Reset all streaming state. Call this between requests to prevent - state leakage when parser instance is reused. - """ - # Reset section state - self._reset_section_state() - - # Reset parent class state - self.current_tool_name_sent = False - self.prev_tool_call_arr = [] - self.current_tool_id = -1 - self.streamed_args_for_tool = [] - - logger.debug("Streaming state reset") + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + request = super().adjust_request(request) + if request.tools and request.tool_choice != "none": + # Ensure special-token markers appear as literal text in + # current_text so we can do pure text-based parsing. + request.skip_special_tokens = False + return request def extract_tool_calls( self, @@ -198,6 +119,95 @@ class KimiK2ToolParser(ToolParser): tools_called=False, tool_calls=[], content=model_output ) + def _extract_content(self, current_text: str) -> str | None: + """Return unsent content before the tool-calls section, or None. + + Holds back any trailing suffix that partially matches + ``<|tool_calls_section_begin|>`` to avoid leaking marker bytes. + """ + if self.tool_calls_start_token not in current_text: + overlap = partial_tag_overlap(current_text, self.tool_calls_start_token) + sendable_idx = len(current_text) - overlap + else: + sendable_idx = current_text.index(self.tool_calls_start_token) + + if sendable_idx > self._sent_content_idx: + content = current_text[self._sent_content_idx : sendable_idx] + self._sent_content_idx = sendable_idx + return content + return None + + def _extract_tool_calls(self, current_text: str) -> list[str]: + """Extract raw bodies from ``<|tool_call_begin|>…<|tool_call_end|>`` blocks.""" + if self.tool_calls_start_token not in current_text: + return [] + + results: list[str] = [] + pos = current_text.index(self.tool_calls_start_token) + while True: + start = current_text.find(self.tool_call_start_token, pos) + if start == -1: + break + tc_start = start + len(self.tool_call_start_token) + end = current_text.find(self.tool_call_end_token, tc_start) + + if end != -1: + tool_call = current_text[tc_start:end] + pos = end + len(self.tool_call_end_token) + else: + tool_call = current_text[tc_start:] + overlap = partial_tag_overlap(tool_call, self.tool_call_end_token) + if overlap: + tool_call = tool_call[:-overlap] + + results.append(tool_call) + + if end == -1: + break + return results + + @staticmethod + def _extract_tool_id_and_name( + header: str | None, + ) -> tuple[str | None, str | None]: + """Parse ``(tool_id, tool_name)`` from a header + like ``"functions.get_weather:0"``.""" + if header is None: + return None, None + match = re.match(r"(.+:\d+)", header) + if not match: + return None, None + + tool_id = match.group(1).strip() + tool_name = tool_id.split(":")[0].split(".")[-1] + return tool_id, tool_name + + def _split_tool_call(self, tool_call: str) -> tuple[str | None, str | None]: + """Split a tool-call body into ``(header, arguments)`` at the argument marker. + + Example:: + 'get_weather:0 <|tool_call_argument_begin|>{"c' + -> ("get_weather:0", '{"c') + """ + arg_pos = tool_call.find(self.tool_call_arg_token) + if arg_pos == -1: + return None, None + header = tool_call[:arg_pos].strip() + tool_args = tool_call[arg_pos + len(self.tool_call_arg_token) :] + return header, tool_args + + def _compute_args_diff(self, index: int, tool_args: str | None) -> str | None: + """Return new argument text not yet sent for tool `index`, or None.""" + if tool_args is None: + return None + prev = self.streamed_args_for_tool[index] + if len(tool_args) <= len(prev): + return None + diff = tool_args[len(prev) :] + self.streamed_args_for_tool[index] = tool_args + self.prev_tool_call_arr[index]["arguments"] = tool_args + return diff + def extract_tool_calls_streaming( self, previous_text: str, @@ -208,394 +218,59 @@ class KimiK2ToolParser(ToolParser): delta_token_ids: Sequence[int], request: ChatCompletionRequest, ) -> DeltaMessage | None: - logger.debug("delta_text: %s", delta_text) - logger.debug("delta_token_ids: %s", delta_token_ids) - - # Flag to defer section exit until after tool parsing completes - deferred_section_exit = False - - # Add delta to buffer for split marker detection - self.token_buffer += delta_text - - # Enforce buffer size limit to prevent memory issues - if len(self.token_buffer) > self.buffer_max_size: - if not self._buffer_overflow_logged: - logger.warning( - "Token buffer exceeded max size (%d bytes), flushing excess. " - "This may indicate very long markers or unusual tokenization.", - self.buffer_max_size, - ) - self._buffer_overflow_logged = True - # Keep only the most recent content that might contain partial markers - self.token_buffer = self.token_buffer[-self.buffer_max_size // 2 :] - - # Check buffer for section markers (handles split tokens) - buffered_text, found_section_begin, found_section_end = ( - self._check_and_strip_markers(self.token_buffer) - ) - - # Track section state transitions - if found_section_begin and not self.in_tool_section: - logger.debug("Entering tool section") - self.in_tool_section = True - self.token_buffer = buffered_text # Use cleaned buffer - self.section_char_count = 0 # Reset counter for new section - - if found_section_end and self.in_tool_section: - logger.debug("Detected section end marker") - # CRITICAL: Don't exit early if tool_call_end is in this chunk. - # Tool parser must emit final arguments/close first to avoid dropping - # the final tool update and leaking tokens into reasoning channel. - has_tool_end = self.tool_call_end_token_id in delta_token_ids - if has_tool_end: - # Defer exit until after tool parsing completes - deferred_section_exit = True - logger.debug("Deferring section exit: tool_call_end in same chunk") - self.token_buffer = buffered_text - else: - # No tool call ending, safe to exit immediately - logger.debug("Exiting tool section") - self._reset_section_state() - # Extract any content AFTER the section end marker in delta_text - # (don't use buffered_text as it contains tool call data) - post_section_content = "" - for variant in self.tool_calls_end_token_variants: - if variant in delta_text: - parts = delta_text.split(variant, 1) - if len(parts) > 1: - post_section_content = parts[1] - break - if post_section_content.strip(): - return DeltaMessage(content=post_section_content) - return DeltaMessage(content="") - else: - self.token_buffer = buffered_text - - # Check if any variant of section start token is in current_token_ids - has_section_token = any( - tid in current_token_ids for tid in self.tool_calls_start_token_ids - ) - - # Early return: if no section token detected yet, return as reasoning content - if not has_section_token and not self.in_tool_section: - logger.debug("No tool call tokens found!") - # Don't clear buffer - it needs to accumulate partial markers across deltas - # Buffer overflow is already protected by lines 215-224 - return DeltaMessage(content=delta_text) - - # Strip section markers from delta_text for subsequent processing - # NOTE: This preprocessing happens BEFORE the regex-based tool call - # parsing (from PR #24847) to ensure markers are removed cleanly - # before pattern matching. No double-stripping occurs because - # section markers and tool call markers are distinct. - delta_text, _, _ = self._check_and_strip_markers(delta_text) - - # Error recovery: If in tool section for too long, force exit - if self.in_tool_section: - self.section_char_count += len(delta_text) - if self.section_char_count > self.max_section_chars: - logger.warning( - "Tool section exceeded max length (%d chars), forcing exit. " - "This may indicate malformed model output.", - self.max_section_chars, - ) - self._reset_section_state() - # Deferred exit already handled by forced exit above - # Return remaining content as reasoning (or empty delta if no content) - return DeltaMessage(content=delta_text if delta_text.strip() else "") - try: - # figure out where we are in the parsing by counting tool call - # start & end tags - prev_tool_start_count = previous_token_ids.count( - self.tool_call_start_token_id - ) - prev_tool_end_count = previous_token_ids.count(self.tool_call_end_token_id) - cur_tool_start_count = current_token_ids.count( - self.tool_call_start_token_id - ) - cur_tool_end_count = current_token_ids.count(self.tool_call_end_token_id) - tool_call_portion = None - text_portion = None + # Extract any content before tool calls. + content = self._extract_content(current_text) + tool_calls = self._extract_tool_calls(current_text) + tool_call_deltas: list[DeltaToolCall] = [] - # case: if we're generating text, OR rounding out a tool call - if ( - cur_tool_start_count == cur_tool_end_count - and prev_tool_end_count == cur_tool_end_count - and self.tool_call_end_token not in delta_text - ): - # Suppress content between section begin and first tool begin - # (header noise). Don't suppress content between tools to avoid - # breaking potential delimiter characters. - if self.in_tool_section and cur_tool_start_count == 0: - logger.debug( - "In tool section before first tool, suppressing: %s", - delta_text, - ) - # Return empty delta to maintain iterator contract - return DeltaMessage(content="") - logger.debug("Generating text content! skipping tool parsing.") - return DeltaMessage(content=delta_text) + for i, tool_call in enumerate(tool_calls): + # First time seeing tool call at index i. + if i >= len(self.prev_tool_call_arr): + # Initialize streaming state. + self.prev_tool_call_arr.append({}) + self.streamed_args_for_tool.append("") - if self.tool_call_end_token in delta_text: - logger.debug("tool_call_end_token in delta_text") - full_text = current_text + delta_text - tool_call_portion = ( - full_text.split(self.tool_call_start_token)[-1] - .split(self.tool_call_end_token)[0] - .rstrip() - ) - delta_text = delta_text.split(self.tool_call_end_token)[0].rstrip() - text_portion = delta_text.split(self.tool_call_end_token)[-1].lstrip() + header, tool_args = self._split_tool_call(tool_call) - # case -- we're starting a new tool call - if ( - cur_tool_start_count > cur_tool_end_count - and cur_tool_start_count > prev_tool_start_count - ): - if len(delta_token_ids) > 1: - tool_call_portion = current_text.split(self.tool_call_start_token)[ - -1 - ] - else: - tool_call_portion = None - delta = None - - text_portion = None - - # set cursors and state appropriately - self.current_tool_id += 1 - self.current_tool_name_sent = False - self.streamed_args_for_tool.append("") - logger.debug("Starting on a new tool %s", self.current_tool_id) - - # case -- we're updating an existing tool call - elif ( - cur_tool_start_count > cur_tool_end_count - and cur_tool_start_count == prev_tool_start_count - ): - # get the portion of the text that's the tool call - tool_call_portion = current_text.split(self.tool_call_start_token)[-1] - text_portion = None - - # case -- the current tool call is being closed. - elif ( - cur_tool_start_count == cur_tool_end_count - and cur_tool_end_count >= prev_tool_end_count - ): - if self.prev_tool_call_arr is None or len(self.prev_tool_call_arr) == 0: - logger.debug("attempting to close tool call, but no tool call") - # Handle deferred section exit before returning - if deferred_section_exit and self.in_tool_section: - self._reset_section_state() - return None - diff = self.prev_tool_call_arr[self.current_tool_id].get("arguments") - if diff: - diff = ( - diff.encode("utf-8").decode("unicode_escape") - if diff is str - else diff - ) - if '"}' not in delta_text: - # Handle deferred section exit before returning - if deferred_section_exit and self.in_tool_section: - self._reset_section_state() - return None - end_loc = delta_text.rindex('"}') - diff = delta_text[:end_loc] + '"}' - logger.debug( - "Finishing tool and found diff that had not " - "been streamed yet: %s", - diff, - ) - self.streamed_args_for_tool[self.current_tool_id] += diff - # Handle deferred section exit before returning - if deferred_section_exit and self.in_tool_section: - logger.debug("Completing deferred section exit") - self._reset_section_state() - return DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.current_tool_id, - function=DeltaFunctionCall(arguments=diff).model_dump( - exclude_none=True - ), - ) - ] - ) - - # case -- otherwise we're just generating text - else: - # Check if we're in tool section - if so, suppress - if self.in_tool_section: - logger.debug("In tool section, suppressing text generation") - # Handle deferred section exit before returning - if deferred_section_exit: - self._reset_section_state() - return DeltaMessage(content="") - text = delta_text.replace(self.tool_call_start_token, "") - text = text.replace(self.tool_call_end_token, "") - delta = DeltaMessage(tool_calls=[], content=text) - # Handle deferred section exit before returning - if deferred_section_exit and self.in_tool_section: - self._reset_section_state() - return delta - - current_tool_call = dict() - if tool_call_portion: - current_tool_call_matches = self.stream_tool_call_portion_regex.match( - tool_call_portion - ) - if current_tool_call_matches: - tool_id, tool_args = current_tool_call_matches.groups() - tool_name = tool_id.split(":")[0].split(".")[-1] - current_tool_call["id"] = tool_id.strip() - current_tool_call["name"] = tool_name - current_tool_call["arguments"] = tool_args - else: - current_tool_call_name_matches = ( - self.stream_tool_call_name_regex.match(tool_call_portion) - ) - if current_tool_call_name_matches: - (tool_id_str,) = current_tool_call_name_matches.groups() - tool_name = tool_id_str.split(":")[0].split(".")[-1] - current_tool_call["id"] = tool_id_str.strip() - current_tool_call["name"] = tool_name - current_tool_call["arguments"] = "" - else: - logger.debug("Not enough token") - return None - - # case - we haven't sent the tool name yet. If it's available, send - # it. otherwise, wait until it's available. - if not self.current_tool_name_sent: - if current_tool_call is None: - return None - function_name: str | None = current_tool_call.get("name") - tool_id = current_tool_call.get("id") - if function_name: - self.current_tool_name_sent = True - return DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.current_tool_id, - type="function", - id=tool_id, - function=DeltaFunctionCall( - name=function_name - ).model_dump(exclude_none=True), - ) - ] - ) - else: - return None - - # case -- otherwise, send the tool call delta - - # if the tool call portion is None, send the delta as text - if tool_call_portion is None: - # if there's text but not tool calls, send that - - # otherwise None to skip chunk - # CRITICAL: Never return content if we're in a tool section - if self.in_tool_section: - return None - delta = ( - DeltaMessage(content=delta_text) - if text_portion is not None - else None - ) - return delta - - # now, the nitty-gritty of tool calls - # now we have the portion to parse as tool call. - - logger.debug( - "Trying to parse current tool call with ID %s", self.current_tool_id - ) - - # if we're starting a new tool call, push an empty object in as - # a placeholder for the arguments - if len(self.prev_tool_call_arr) <= self.current_tool_id: - self.prev_tool_call_arr.append({}) - - # main logic for tool parsing here - compare prev. partially-parsed - # JSON to the current partially-parsed JSON - prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get( - "arguments" - ) - cur_arguments = current_tool_call.get("arguments") - - logger.debug("diffing old arguments: %s", prev_arguments) - logger.debug("against new ones: %s", cur_arguments) - - # case -- no arguments have been created yet. skip sending a delta. - if not cur_arguments and not prev_arguments: - logger.debug("Skipping text %s - no arguments", delta_text) - delta = None - - # case -- prev arguments are defined, but non are now. - # probably impossible, but not a fatal error - just keep going - elif not cur_arguments and prev_arguments: - logger.error( - "should be impossible to have arguments reset " - "mid-call. skipping streaming anything." - ) - delta = None - - # case -- we now have the first info about arguments available from - # autocompleting the JSON - elif cur_arguments and not prev_arguments: - delta = DeltaMessage( - tool_calls=[ + # Stream back tool name. + if "name" not in self.prev_tool_call_arr[i]: + tool_id, tool_name = self._extract_tool_id_and_name(header) + if not tool_name: + # Can't skip to tool i+1 if i isn't ready + break + self.prev_tool_call_arr[i]["name"] = tool_name + self.prev_tool_call_arr[i]["id"] = tool_id + tool_call_deltas.append( DeltaToolCall( - index=self.current_tool_id, - function=DeltaFunctionCall( - arguments=cur_arguments - ).model_dump(exclude_none=True), + index=i, + type="function", + id=tool_id, + function=DeltaFunctionCall(name=tool_name).model_dump( + exclude_none=True + ), ) - ] - ) - self.streamed_args_for_tool[self.current_tool_id] = cur_arguments - - # last case -- we have an update to existing arguments. - elif cur_arguments and prev_arguments: - if ( - isinstance(delta_text, str) - and cur_arguments != prev_arguments - and len(cur_arguments) > len(prev_arguments) - and cur_arguments.startswith(prev_arguments) - ): - delta_arguments = cur_arguments[len(prev_arguments) :] - logger.debug("got diff %s", delta_text) - - delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.current_tool_id, - function=DeltaFunctionCall( - arguments=delta_arguments - ).model_dump(exclude_none=True), - ) - ] ) - self.streamed_args_for_tool[self.current_tool_id] = cur_arguments - else: - delta = None - # handle saving the state for the current tool into - # the "prev" list for use in diffing for the next iteration - if self.current_tool_id == len(self.prev_tool_call_arr) - 1: - self.prev_tool_call_arr[self.current_tool_id] = current_tool_call - else: - self.prev_tool_call_arr.append(current_tool_call) + # Stream back new tool args by diffing against what was sent. + args_diff = self._compute_args_diff(i, tool_args) + if args_diff: + tool_call_deltas.append( + DeltaToolCall( + index=i, + function=DeltaFunctionCall(arguments=args_diff).model_dump( + exclude_none=True + ), + ) + ) - # Handle deferred section exit after tool parsing completes - if deferred_section_exit and self.in_tool_section: - logger.debug("Completing deferred section exit") - self._reset_section_state() - - return delta + if content or tool_call_deltas: + return DeltaMessage( + content=content, + tool_calls=tool_call_deltas, + ) + return None except Exception: logger.exception("Error trying to handle streaming tool call.") - return None # do not stream a delta. skip this token ID. + return None From 45232a454e4c57191633fefda5baa9fa1f5dc8fd Mon Sep 17 00:00:00 2001 From: TJian Date: Sun, 19 Apr 2026 17:57:39 +0800 Subject: [PATCH 117/150] [FEAT] [Perf] [Gemma4] Fused Gemma4 Routing Function Triton (#39083) Signed-off-by: tjtanaa --- pyproject.toml | 1 + tests/kernels/moe/test_gemma4router.py | 57 +++++++++++ vllm/model_executor/models/gemma4.py | 136 ++++++++++++++++++++++--- 3 files changed, 179 insertions(+), 15 deletions(-) create mode 100644 tests/kernels/moe/test_gemma4router.py diff --git a/pyproject.toml b/pyproject.toml index f55dd9308bd..5c87de018c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -170,6 +170,7 @@ eles = "eles" datas = "datas" ser = "ser" ure = "ure" +VALU = "VALU" # Walsh-Hadamard Transform wht = "wht" WHT = "WHT" diff --git a/tests/kernels/moe/test_gemma4router.py b/tests/kernels/moe/test_gemma4router.py new file mode 100644 index 00000000000..ba69d692749 --- /dev/null +++ b/tests/kernels/moe/test_gemma4router.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from vllm.model_executor.models.gemma4 import ( + gemma4_fused_routing_kernel_triton, + gemma4_routing_function_torch, +) + + +def sort_by_id(w, ids): + order = ids.argsort(dim=-1) + return w.gather(1, order), ids.gather(1, order) + + +# Gemma4 MoE Model has context length of 250K +# the minus 1 is to ensure that edge cases are tested +@pytest.mark.parametrize("num_tokens", [1, 2, 2048, 250000]) +@pytest.mark.parametrize("num_experts", [128]) # gemma4 moe experts +@pytest.mark.parametrize("topk", [8]) # gemma4 topk +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32]) +def test_gemma4_routing_kernel_triton( + num_tokens: int, + num_experts: int, + topk: int, + dtype: torch.dtype, +): + torch.manual_seed(0) + + gating = torch.randn(num_tokens, num_experts, dtype=dtype, device="cuda") + scales = torch.rand(num_experts, dtype=torch.float32, device="cuda") + + ref_w, ref_ids = gemma4_routing_function_torch(gating, topk, scales) + tri_w, tri_ids = gemma4_fused_routing_kernel_triton(gating, topk, scales) + + # Sort by expert id — to remove tie-breaking differences + ref_ws, ref_is = sort_by_id(ref_w, ref_ids) + tri_ws, tri_is = sort_by_id(tri_w, tri_ids) + + ids_match = (ref_is == tri_is).all().item() + weights_match = torch.allclose(ref_ws, tri_ws, atol=1e-2, rtol=1e-2) + all_match = ids_match and weights_match + max_err = (ref_ws - tri_ws).abs().max().item() + print( + f"T={num_tokens:5d} E={num_experts:4d} K={topk} " + f"{str(dtype).split('.')[-1]:7s} ids={ids_match} max_Δweight={max_err:.2e}" + ) + if not all_match: + bad = (ref_is != tri_is).any(dim=-1).nonzero(as_tuple=True)[0] + if len(bad): + r = bad[0].item() + print( + f" first bad row {r}: ref_ids={ref_ids[r].tolist()} " + f"tri_ids={tri_ids[r].tolist()}" + ) + assert all_match diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 06189540090..d166a9df38a 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -57,7 +57,9 @@ from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata from .interfaces import ( @@ -79,6 +81,120 @@ from .utils import ( logger = init_logger(__name__) +@triton.jit +def _gemma4_routing_kernel( + gating_ptr, + per_expert_scale_ptr, + topk_weights_ptr, + topk_ids_ptr, + E: tl.constexpr, + K: tl.constexpr, + BLOCK_E: tl.constexpr, +): + pid = tl.program_id(0) + offs_e = tl.arange(0, BLOCK_E) + valid = offs_e < E + + logits = tl.load( + gating_ptr + pid * E + offs_e, + mask=valid, + other=-float("inf"), + ).to(tl.float32) + + max_l = tl.max(logits, axis=0) + + # Float32 → ascending-sortable bijection + MIN32 = -2147483648 + logit_bits = logits.to(tl.int32, bitcast=True) + sign_b = logit_bits >> 31 + key = tl.where(sign_b == 0, logit_bits ^ -1, logit_bits ^ MIN32) + key = tl.where(valid, key, 0x7FFFFFFF) + sk64 = key.to(tl.int64) & 0x00000000FFFFFFFF + packed = (sk64 << 32) | offs_e.to(tl.int64) + sorted_p = tl.sort(packed, descending=False) + + # Vectorized extraction of ALL sorted elements — no K-loop, no cross-lane reductions + all_keys = ((sorted_p >> 32) & 0x00000000FFFFFFFF).to(tl.int32) + all_ids = (sorted_p & 0x00000000FFFFFFFF).to(tl.int32) + + # Inverse bijection: recover original logit bits + sign_k = all_keys >> 31 + all_bits = tl.where(sign_k < 0, all_keys ^ -1, all_keys ^ MIN32) + all_logits = all_bits.to(tl.float32, bitcast=True) + + # Compute raw_exp for ALL BLOCK_E elements — vectorized, ~2 VALU clocks + all_raw_exp = tl.math.exp2((all_logits - max_l) * 1.4426950408889634) + + # Sum only top-K for renorm — ONE masked reduction + top_mask = offs_e < K + renorm_raw = tl.sum(tl.where(top_mask, all_raw_exp, 0.0), axis=0) + renorm_raw = tl.where(renorm_raw > 0.0, renorm_raw, 1.0) + inv_renorm = 1.0 / renorm_raw + + # Load scales for top-K only (masked gather; scale array is tiny → L1 cached) + all_scales = tl.load( + per_expert_scale_ptr + all_ids.to(tl.int64), + mask=top_mask, + other=1.0, + ).to(tl.float32) + + # Final weights: vectorized multiply (only top-K will be stored) + all_weights = (all_raw_exp * inv_renorm * all_scales).to(tl.float32) + + # Write results with TWO masked stores — replaces K × 2 serial scalar stores + base_off = pid * K + offs_e + tl.store(topk_ids_ptr + base_off, all_ids, mask=top_mask) + tl.store(topk_weights_ptr + base_off, all_weights, mask=top_mask) + + +def gemma4_fused_routing_kernel_triton( + gating_output: torch.Tensor, + topk: int, + per_expert_scale: torch.Tensor, + num_warps: int = 1, +) -> tuple[torch.Tensor, torch.Tensor]: + gating_output = gating_output.contiguous() + per_expert_scale = per_expert_scale.contiguous() + T, E = gating_output.shape + weights = torch.empty(T, topk, dtype=torch.float32, device=gating_output.device) + ids = torch.empty(T, topk, dtype=torch.int32, device=gating_output.device) + BLOCK_E = triton.next_power_of_2(E) + _gemma4_routing_kernel[(T,)]( + gating_output, + per_expert_scale, + weights, + ids, + E, + topk, + BLOCK_E, + num_warps=num_warps, + ) + return weights, ids + + +def gemma4_routing_function_torch( + gating_output: torch.Tensor, + topk: int, + per_expert_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + _, topk_ids = torch.topk(gating_output, k=topk, dim=-1) + router_probabilities = torch.nn.functional.softmax(gating_output, dim=-1) + indicator = torch.nn.functional.one_hot( + topk_ids, num_classes=gating_output.size(-1) + ).sum(dim=-2) + gate_weights = indicator * router_probabilities + renorm_factor = torch.sum(gate_weights, dim=-1, keepdim=True) + renorm_factor = torch.where(renorm_factor > 0.0, renorm_factor, 1.0) + dispatch_weights = gate_weights / renorm_factor + + topk_weights = dispatch_weights.gather(1, topk_ids) + + # Fold per_expert_scale into routing weights + expert_scales = per_expert_scale[topk_ids].to(topk_weights.dtype) + topk_weights = topk_weights * expert_scales + return topk_weights.to(torch.float32), topk_ids.to(torch.int32) + + def _get_text_config(config): """Dereference text_config if config is a nested Gemma4Config. @@ -216,22 +332,12 @@ class Gemma4MoE(nn.Module): topk: int, renormalize: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - _, topk_ids = torch.topk(gating_output, k=topk, dim=-1) - router_probabilities = torch.nn.functional.softmax(gating_output, dim=-1) - indicator = torch.nn.functional.one_hot( - topk_ids, num_classes=gating_output.size(-1) - ).sum(dim=-2) - gate_weights = indicator * router_probabilities - renorm_factor = torch.sum(gate_weights, dim=-1, keepdim=True) - renorm_factor = torch.where(renorm_factor > 0.0, renorm_factor, 1.0) - dispatch_weights = gate_weights / renorm_factor + if current_platform.is_cuda_alike() or current_platform.is_xpu(): + return gemma4_fused_routing_kernel_triton( + gating_output, topk, per_expert_scale + ) - topk_weights = dispatch_weights.gather(1, topk_ids) - - # Fold per_expert_scale into routing weights - expert_scales = per_expert_scale[topk_ids].to(topk_weights.dtype) - topk_weights = topk_weights * expert_scales - return topk_weights.to(torch.float32), topk_ids.to(torch.int32) + return gemma4_routing_function_torch(gating_output, topk, per_expert_scale) # FusedMoE experts with custom Gemma4 routing self.experts = FusedMoE( From 982beae809b1690f745d9d52616de4d1da2c5855 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Sun, 19 Apr 2026 18:06:20 +0300 Subject: [PATCH 118/150] Optimize nemotron VL image/video preprocessing (#40283) Signed-off-by: milesial Co-authored-by: milesial --- .../processors/nano_nemotron_vl.py | 216 ++++++++++-------- 1 file changed, 117 insertions(+), 99 deletions(-) diff --git a/vllm/transformers_utils/processors/nano_nemotron_vl.py b/vllm/transformers_utils/processors/nano_nemotron_vl.py index 57e874be4ec..76b73d21635 100644 --- a/vllm/transformers_utils/processors/nano_nemotron_vl.py +++ b/vllm/transformers_utils/processors/nano_nemotron_vl.py @@ -8,7 +8,6 @@ # -------------------------------------------------------- import math -import warnings from abc import ABC, abstractmethod from collections.abc import Sequence from dataclasses import dataclass @@ -26,7 +25,7 @@ from transformers import BatchFeature, PretrainedConfig, TensorType from vllm.model_executor.models.parakeet import ParakeetExtractor from vllm.multimodal.evs import compute_retained_tokens_count from vllm.multimodal.inputs import AudioItem -from vllm.multimodal.processing.processor import PromptUpdateDetails, _seq2tokens +from vllm.multimodal.processing.processor import PromptUpdateDetails from vllm.tokenizers.hf import HfTokenizer from .internvl import calculate_internvl_targets, get_internvl_target_ratios @@ -63,42 +62,50 @@ def calculate_timestamps( return timestamps -def input_conditioner(x: torch.Tensor, norm_mean: torch.Tensor, norm_std: torch.Tensor): - return (x - norm_mean) / norm_std - - -def _bicubic_from_ndarray( - array: npt.NDArray[Any], *, size: tuple[int, int] +@torch.compile(dynamic=True) +def _bicubic_resize_and_normalize( + tensor: torch.Tensor, + size: tuple[int, int] | None = None, + norm_mean: torch.Tensor | None = None, + norm_std: torch.Tensor | None = None, + dtype: torch.dtype = torch.float32, ) -> torch.Tensor: - """ - Convert a 4D NHWC ndarray to NCHW and interpolate with bicubic. - Suppresses PyTorch's non-writable NumPy warning because interpolate copies, - and torch.from_numpy(array) is discarded at the end of function scope. - """ + """Permute NHWC→NCHW, optional bicubic resize, rescale + normalize. - with warnings.catch_warnings(): - msg = "The given NumPy array is not writ.*" - # Apparently, different versions of PyTorch use writable or writeable. - warnings.filterwarnings("ignore", message=msg, category=UserWarning) - tensor = torch.from_numpy(array) - assert tensor.ndim == 4, f"{tensor.ndim=}" - tensor = tensor.permute(0, 3, 1, 2) - return ( - torch.nn.functional.interpolate( + Input must be a raw 4-D **NHWC** tensor. + + *size*: target ``(H, W)``; skips interpolation when ``None``. + *norm_mean* / *norm_std*: when both provided, fused + ``(x/255 - mean) / std`` + dtype cast; otherwise ``x/255`` + cast. + """ + tensor = tensor.permute(0, 3, 1, 2).to(dtype=torch.float32) + if size is not None: + tensor = torch.nn.functional.interpolate( tensor, size=size, mode="bicubic", align_corners=False, antialias=True ) - / 255.0 + if norm_mean is not None and norm_std is not None: + return ((tensor / 255.0 - norm_mean) / norm_std).to(dtype=dtype).contiguous() + return (tensor / 255.0).to(dtype=dtype).contiguous() + + +def _pil_to_nhwc_tensor(image: Image.Image) -> torch.Tensor: + """Convert a PIL image to a 4-D NHWC tensor suitable for compiled ops.""" + array = np.asarray( + image.convert("RGB") if image.mode != "RGB" else image, dtype=np.uint8 ) + return torch.from_numpy(np.expand_dims(array, axis=0)) def dynamic_preprocess( - image, + image: Image.Image, *, - image_size=512, - max_num_tiles=12, - use_thumbnail=True, - idx=0, -): + image_size: int = 512, + max_num_tiles: int = 12, + use_thumbnail: bool = True, + norm_mean: torch.Tensor | None = None, + norm_std: torch.Tensor | None = None, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: orig_width, orig_height = image.size target_ratios = get_internvl_target_ratios(1, max_num_tiles) @@ -111,13 +118,15 @@ def dynamic_preprocess( use_thumbnail=False, ) - image = np.asarray( - image.convert("RGB") if image.mode != "RGB" else image, dtype=np.uint8 + tensor = _pil_to_nhwc_tensor(image) + + resized_img = _bicubic_resize_and_normalize( + tensor, + size=(target_height, target_width), + norm_mean=norm_mean, + norm_std=norm_std, + dtype=dtype, ) - - image = np.expand_dims(image, axis=0) - - resized_img = _bicubic_from_ndarray(image, size=(target_height, target_width)) B, C, H, W = resized_img.shape hp, wp = H // image_size, W // image_size patches = ( @@ -127,30 +136,16 @@ def dynamic_preprocess( ) if use_thumbnail and patches.shape[0] > 1: - thumb = _bicubic_from_ndarray(image, size=(image_size, image_size)) + thumb = _bicubic_resize_and_normalize( + tensor, + size=(image_size, image_size), + norm_mean=norm_mean, + norm_std=norm_std, + dtype=dtype, + ) patches = torch.cat([patches, thumb], dim=0) - return list(patches) - - -def image_to_pixel_values( - image: Image.Image, - *, - input_size: int, - max_num: int, - use_thumbnail: bool, - idx: int, -) -> torch.Tensor: - images = dynamic_preprocess( - image, - image_size=input_size, - max_num_tiles=max_num, - use_thumbnail=use_thumbnail, - idx=idx, - ) - - pixel_values = torch.stack(images) - return pixel_values + return patches def _compute_aspect_preserving_size( @@ -233,14 +228,16 @@ def video_to_pixel_values( video_maintain_aspect_ratio: bool = False, patch_size: int = 16, downsample_ratio: float = 0.5, + norm_mean: torch.Tensor | None = None, + norm_std: torch.Tensor | None = None, + dtype: torch.dtype = torch.float32, ) -> torch.Tensor: - # (num_frames, H, W, C) -> (num_frames, C, H, W) - video_tensor = torch.from_numpy(video).permute(0, 3, 1, 2) + """Convert video ndarray (T, H, W, C) to normalized pixel tensor (T, C, H, W).""" + orig_h, orig_w = video.shape[1], video.shape[2] + size: tuple[int, int] | None = None if video_target_num_patches is not None: - # Resize to target patch count (aspect-preserving or square). - orig_h, orig_w = video_tensor.shape[2], video_tensor.shape[3] - target_w, target_h, _ = get_video_target_size_and_feature_size( + tw, th, _ = get_video_target_size_and_feature_size( orig_w=orig_w, orig_h=orig_h, target_patches=video_target_num_patches, @@ -248,14 +245,13 @@ def video_to_pixel_values( patch_size=patch_size, downsample_ratio=downsample_ratio, ) - if video_tensor.shape[2] != target_h or video_tensor.shape[3] != target_w: - return _bicubic_from_ndarray(video, size=(target_h, target_w)) - elif video_tensor.shape[2] != input_size or video_tensor.shape[3] != input_size: - return _bicubic_from_ndarray(video, size=(input_size, input_size)) + if orig_h != th or orig_w != tw: + size = (th, tw) + elif orig_h != input_size or orig_w != input_size: + size = (input_size, input_size) - video_tensor = video_tensor / 255.0 - - return video_tensor + tensor = torch.from_numpy(video) + return _bicubic_resize_and_normalize(tensor, size, norm_mean, norm_std, dtype) class DynamicResolutionImageTiler: @@ -343,6 +339,7 @@ class DynamicResolutionImageTiler: self, text_prompt_length: int, images: list[Image.Image], + dtype: torch.dtype = torch.float32, ) -> tuple[list[torch.Tensor], list[int]]: num_tokens_available = self.max_num_tokens_available(text_prompt_length) params_per_image = self.compute_params(images, num_tokens_available) @@ -350,7 +347,7 @@ class DynamicResolutionImageTiler: feature_sizes = [] images = [] for param in params_per_image: - for t in self.apply_params(param): + for t in self.apply_params(param, dtype=dtype): assert t.ndim == 3, f"{t.ndim=}: expected 3 dim tensor" images.append(t) feature_sizes.append(param.num_embeddings) @@ -363,17 +360,23 @@ class DynamicResolutionImageTiler: num_embeddings: int patch_size: tuple[int, int] - def apply_params(self, params: DynamicResolutionParams) -> list[torch.Tensor]: + def apply_params( + self, + params: DynamicResolutionParams, + dtype: torch.dtype = torch.float32, + ) -> list[torch.Tensor]: target_size = ( params.patch_size[1] * self._patch_size, params.patch_size[0] * self._patch_size, ) - image = np.asarray( - params.media.convert("RGB") if params.media.mode != "RGB" else params.media, - dtype=np.uint8, + tensor = _pil_to_nhwc_tensor(params.media) + resized_img = _bicubic_resize_and_normalize( + tensor, + size=target_size, + norm_mean=self.norm_mean, + norm_std=self.norm_std, + dtype=dtype, ) - image = np.expand_dims(image, axis=0) - resized_img = _bicubic_from_ndarray(image, size=target_size) return list(resized_img) def process_media( @@ -619,6 +622,7 @@ class BaseNanoNemotronVLProcessor(ABC): norm_mean=config.norm_mean, norm_std=config.norm_std, ) + self.dtype: torch.dtype = getattr(config, "dtype", torch.float32) @staticmethod def use_dynamic_resolution(config: PretrainedConfig) -> bool: @@ -662,14 +666,16 @@ class BaseNanoNemotronVLProcessor(ABC): max_num_tiles: int, ) -> list[torch.Tensor]: return [ - image_to_pixel_values( + dynamic_preprocess( image, - input_size=self.image_size, - max_num=max_num_tiles, + image_size=self.image_size, + max_num_tiles=max_num_tiles, use_thumbnail=self.use_thumbnail, - idx=idx, + norm_mean=self.norm_mean, + norm_std=self.norm_std, + dtype=self.dtype, ) - for idx, image in enumerate(images) + for image in images ] def _preprocess_image( @@ -690,23 +696,22 @@ class BaseNanoNemotronVLProcessor(ABC): pixel_values_lst, num_tokens_per_image = tiler._images_to_pixel_values_lst( text_prompt_length=text_prompt_length, images=images, + dtype=self.dtype, ) imgs_sizes = [(pv.shape[-2], pv.shape[-1]) for pv in pixel_values_lst] - normalized = [ - input_conditioner(img, tiler.norm_mean, tiler.norm_std) - for img in pixel_values_lst - ] image_num_patches = torch.tensor([1] * len(num_tokens_per_image)) image_inputs = { - "pixel_values_flat": normalized, + "pixel_values_flat": pixel_values_lst, "imgs_sizes": imgs_sizes, "num_tokens_per_image": num_tokens_per_image, } else: pixel_values_lst = self._images_to_pixel_values_lst(images, max_num_tiles) image_num_patches = torch.tensor([len(item) for item in pixel_values_lst]) - pixel_values_flat = input_conditioner( - torch.cat(pixel_values_lst), self.norm_mean, self.norm_std + pixel_values_flat = ( + torch.cat(pixel_values_lst) + if len(pixel_values_lst) > 1 + else pixel_values_lst[0] ) image_inputs = { "pixel_values_flat": pixel_values_flat, @@ -863,6 +868,8 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): def _videos_to_pixel_values_lst( self, videos: list[npt.NDArray], + *, + dtype: torch.dtype = torch.float32, ) -> list[torch.Tensor]: return [ video_to_pixel_values( @@ -872,6 +879,9 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): video_maintain_aspect_ratio=self.video_maintain_aspect_ratio, patch_size=self.config.patch_size, downsample_ratio=self.config.downsample_ratio, + norm_mean=self.norm_mean, + norm_std=self.norm_std, + dtype=dtype, ) for video in videos ] @@ -886,8 +896,10 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): videos_lst = [v[0] for v in videos] video_metadata_lst = [v[1] for v in videos] + pixel_values_lst_video = self._videos_to_pixel_values_lst( videos_lst, + dtype=self.dtype, ) # We use frame duration in milliseconds (as integer) to ensure @@ -903,10 +915,15 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): metadata["frames_indices"] for metadata in video_metadata_lst ] video_num_patches = torch.tensor([len(item) for item in pixel_values_lst_video]) + + # Normalization already fused into resize above. + # Skip the torch.cat copy when there is exactly one video + if len(pixel_values_lst_video) == 1: + pixel_values_flat = pixel_values_lst_video[0] + else: + pixel_values_flat = torch.cat(pixel_values_lst_video) video_inputs = { - "pixel_values_flat_video": input_conditioner( - torch.cat(pixel_values_lst_video), self.norm_mean, self.norm_std - ), + "pixel_values_flat_video": pixel_values_flat, "video_num_patches": video_num_patches, "frames_indices": frames_indices_lst, "frame_duration_ms": torch.tensor(frame_duration_ms_lst), @@ -1168,20 +1185,21 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): for i, _ in enumerate(tokens_per_frame) ] - # Tokenize frame separator independently - frame_separators_tokenized = [ - _seq2tokens(tokenizer, sep) for sep in frame_separators - ] + # Batch-tokenize all frame separators at once — the HuggingFace + # tokenizers Rust backend parallelizes batch encoding across threads. + batch_encoded = tokenizer( + frame_separators, + add_special_tokens=False, + return_attention_mask=False, + ) + frame_separators_tokenized: list[list[int]] = batch_encoded["input_ids"] # Tokenize each component independently to avoid tokenizer merging tokens # across boundaries. This ensures consistent tokenization regardless of # num_tokens_per_frame values. all_token_ids = [] for i, num_tokens in enumerate(tokens_per_frame): - frame_sep_token_ids = frame_separators_tokenized[i] - all_token_ids.extend(frame_sep_token_ids) - - # Add pre-tokenized special tokens + all_token_ids.extend(frame_separators_tokenized[i]) all_token_ids.extend(img_start_token_ids) all_token_ids.extend(img_context_token_ids * num_tokens) all_token_ids.extend(img_end_token_ids) From d1135a508705cc899946eae2d0cb63b7d9e94bfe Mon Sep 17 00:00:00 2001 From: danisereb Date: Sun, 19 Apr 2026 20:18:40 +0300 Subject: [PATCH 119/150] Fix MoE backend selection for LoRA (unquantized MoE) (#40273) Signed-off-by: Daniel Serebrenik --- .../moe/test_unquantized_backend_selection.py | 85 +++++++++++++++++++ .../layers/fused_moe/oracle/unquantized.py | 5 ++ 2 files changed, 90 insertions(+) diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index 469c45b2312..73aa54fa579 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -11,6 +11,11 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( ) from vllm.platforms import current_platform +skipif_not_cuda_rocm = pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_rocm()), + reason="Only supported on CUDA/ROCm platforms.", +) + @pytest.mark.parametrize( "platform_method,expected_backend", @@ -190,3 +195,83 @@ def test_select_cuda_flashinfer_cutlass_backend( assert selected_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS assert experts_cls is not None + + +@skipif_not_cuda_rocm +def test_select_lora_backend_prefers_triton(): + """LoRA-enabled unquantized MoE should select Triton backend.""" + moe_config = make_dummy_moe_config() + moe_config.is_lora_enabled = True + selected_backend, experts_cls = select_unquantized_moe_backend( + moe_config=moe_config + ) + + assert selected_backend == UnquantizedMoeBackend.TRITON + assert experts_cls is not None + + +@skipif_not_cuda_rocm +def test_select_lora_explicit_non_triton_backend(): + """LoRA should override explicit non-Triton backend to Triton.""" + moe_config = make_dummy_moe_config() + moe_config.is_lora_enabled = True + + # Use string from mapping in function map_unquantized_backend() + moe_config.moe_backend = "flashinfer_cutlass" + + selected_backend, experts_cls = select_unquantized_moe_backend( + moe_config=moe_config + ) + + assert selected_backend == UnquantizedMoeBackend.TRITON + assert experts_cls is not None + + +@skipif_not_cuda_rocm +@pytest.mark.parametrize("is_lora_enabled", [False, True]) +def test_select_explicit_triton_backend(is_lora_enabled): + """Explicit triton backend selection should return Triton.""" + moe_config = make_dummy_moe_config() + moe_config.is_lora_enabled = is_lora_enabled + moe_config.moe_backend = "triton" + + selected_backend, experts_cls = select_unquantized_moe_backend( + moe_config=moe_config + ) + + assert selected_backend == UnquantizedMoeBackend.TRITON + assert experts_cls is not None + + +@skipif_not_cuda_rocm +def test_select_explicit_triton_ignores_flashinfer_env(monkeypatch): + """Explicit triton backend should override FlashInfer env selection.""" + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1") + monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") + + moe_config = make_dummy_moe_config() + moe_config.is_lora_enabled = False + moe_config.moe_backend = "triton" + + selected_backend, experts_cls = select_unquantized_moe_backend( + moe_config=moe_config + ) + + assert selected_backend == UnquantizedMoeBackend.TRITON + assert experts_cls is not None + + +@skipif_not_cuda_rocm +def test_select_lora_ignores_flashinfer_env(monkeypatch): + """LoRA path should still choose Triton even if FlashInfer env is on.""" + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1") + monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") + + moe_config = make_dummy_moe_config() + moe_config.is_lora_enabled = True + selected_backend, experts_cls = select_unquantized_moe_backend( + moe_config=moe_config + ) + + assert selected_backend == UnquantizedMoeBackend.TRITON + assert experts_cls is not None diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 33bf7a0c75f..8fcb8fa1da1 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -163,6 +163,11 @@ def select_unquantized_moe_backend( if current_platform.is_out_of_tree(): return UnquantizedMoeBackend.OOT, None + if moe_config.is_lora_enabled: + return UnquantizedMoeBackend.TRITON, backend_to_kernel_cls( + UnquantizedMoeBackend.TRITON + ) + # NOTE: the kernels are selected in the following order. AVAILABLE_BACKENDS = _get_priority_backends(moe_config) From f150107efd15a873617a65b09189c945b27d61f5 Mon Sep 17 00:00:00 2001 From: Andrew Barnes Date: Sun, 19 Apr 2026 14:34:33 -0400 Subject: [PATCH 120/150] [ROCm] Fix cu_seqlens_q off-by-one in AITER FA speculative decode path (#39120) Signed-off-by: Bortlesboat --- vllm/v1/attention/backends/rocm_aiter_fa.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 14c473a70ab..a6826045ef0 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -1181,7 +1181,7 @@ class AiterFlashAttentionImpl(AttentionImpl): ) descale_shape = ( - attn_metadata.query_start_loc[:num_decodes].shape[0] - 1, + num_decodes, key_cache.shape[2], ) unified_attention( @@ -1189,7 +1189,7 @@ class AiterFlashAttentionImpl(AttentionImpl): k=key_cache, v=value_cache, out=output[:num_decode_tokens], - cu_seqlens_q=attn_metadata.query_start_loc[:num_decodes], + cu_seqlens_q=attn_metadata.query_start_loc[: num_decodes + 1], max_seqlen_q=decode_max_query_len, seqused_k=attn_metadata.seq_lens[:num_decodes], max_seqlen_k=attn_metadata.max_seq_len, From 629d45eacb2c225ccabeb5410c0b4f54eeed2eef Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sun, 19 Apr 2026 15:37:53 -0700 Subject: [PATCH 121/150] [ci] Make ecr authenticate non blocking (#40305) Signed-off-by: Kevin H. Luu --- .buildkite/image_build/image_build.sh | 4 ++-- .buildkite/image_build/image_build_cpu.sh | 2 +- .buildkite/image_build/image_build_cpu_arm64.sh | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.buildkite/image_build/image_build.sh b/.buildkite/image_build/image_build.sh index 9131dfc71a0..00ae34bba6d 100755 --- a/.buildkite/image_build/image_build.sh +++ b/.buildkite/image_build/image_build.sh @@ -92,8 +92,8 @@ check_and_skip_if_image_exists() { } ecr_login() { - aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" - aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com + aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true + aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true } prepare_cache_tags() { diff --git a/.buildkite/image_build/image_build_cpu.sh b/.buildkite/image_build/image_build_cpu.sh index ccfe155fa2b..035f070ab89 100755 --- a/.buildkite/image_build/image_build_cpu.sh +++ b/.buildkite/image_build/image_build_cpu.sh @@ -11,7 +11,7 @@ REPO=$2 BUILDKITE_COMMIT=$3 # authenticate with AWS ECR -aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" +aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true # skip build if image already exists if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then diff --git a/.buildkite/image_build/image_build_cpu_arm64.sh b/.buildkite/image_build/image_build_cpu_arm64.sh index ff3d11c8d59..b561e2c2e46 100755 --- a/.buildkite/image_build/image_build_cpu_arm64.sh +++ b/.buildkite/image_build/image_build_cpu_arm64.sh @@ -11,7 +11,7 @@ REPO=$2 BUILDKITE_COMMIT=$3 # authenticate with AWS ECR -aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" +aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true # skip build if image already exists if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu) ]]; then From 898beca5a8f173e458b4a1f0bac0b2fead0a35a7 Mon Sep 17 00:00:00 2001 From: Liangliang Ma Date: Mon, 20 Apr 2026 08:24:50 +0800 Subject: [PATCH 122/150] [BugFix][XPU] fix lora ops bgmv_expand size not match (#39989) Signed-off-by: Ma, Liangliang Co-authored-by: Kunshang Ji --- vllm/lora/ops/xpu_ops/lora_ops.py | 39 ++++++++++++++++++++++++++++--- 1 file changed, 36 insertions(+), 3 deletions(-) diff --git a/vllm/lora/ops/xpu_ops/lora_ops.py b/vllm/lora/ops/xpu_ops/lora_ops.py index 6d1751c3738..070fd864582 100644 --- a/vllm/lora/ops/xpu_ops/lora_ops.py +++ b/vllm/lora/ops/xpu_ops/lora_ops.py @@ -27,9 +27,42 @@ def bgmv_expand( lora_indices_tensor: torch.Tensor, add_inputs: bool = True, ) -> None: - torch.ops._xpu_C.bgmv_expand( - output_tensor, inputs, lora_b_weights, lora_indices_tensor, add_inputs - ) + weight_out_dim = lora_b_weights.size(-2) + output_dim = output_tensor.size(1) + + if weight_out_dim == output_dim: + torch.ops._xpu_C.bgmv_expand( + output_tensor, + inputs, + lora_b_weights, + lora_indices_tensor, + add_inputs, + ) + elif weight_out_dim < output_dim: + # LoRA weight output dim can be smaller than the output tensor + # (e.g. vocab_size vs padded logits). Use expand_slice to write + # only the matching portion, mirroring torch_ops common_len logic. + torch.ops._xpu_C.bgmv_expand_slice( + output_tensor, + inputs, + lora_b_weights, + lora_indices_tensor, + 0, + weight_out_dim, + add_inputs, + ) + else: + # Weight output dim larger than output tensor: truncate weights. + lora_b_weights = lora_b_weights[..., :output_dim, :].contiguous() + torch.ops._xpu_C.bgmv_expand_slice( + output_tensor, + inputs, + lora_b_weights, + lora_indices_tensor, + 0, + output_dim, + add_inputs, + ) def bgmv_expand_slice( From d886c26d4d4fef7d079696beb4ece1cfb4b008a8 Mon Sep 17 00:00:00 2001 From: Lxx Date: Sun, 19 Apr 2026 19:27:32 -0700 Subject: [PATCH 123/150] [Doc] Fix typos in token_embed pooling documentation (#40266) Signed-off-by: YifanLi3 --- docs/models/pooling_models/token_embed.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/models/pooling_models/token_embed.md b/docs/models/pooling_models/token_embed.md index b0e094267db..c20ae7cf164 100644 --- a/docs/models/pooling_models/token_embed.md +++ b/docs/models/pooling_models/token_embed.md @@ -9,14 +9,14 @@ - Online APIs: - Pooling API (`/pooling`) -The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs a embedding for each token. +The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs an embedding for each token. Many embedding models support both (sequence) embedding and token embedding. For further details on (sequence) embedding, please refer to [this page](embed.md). !!! note Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (embed) is not - what you want, you need to manually specify it via via `PoolerConfig(task="token_embed")` offline or + what you want, you need to manually specify it via `PoolerConfig(task="token_embed")` offline or `--pooler-config.task token_embed` online. ## Typical Use Cases From fcb31c1ac3f5af4701d45007cb2b66d084328f39 Mon Sep 17 00:00:00 2001 From: Shinichi Hemmi <50256998+Alnusjaponica@users.noreply.github.com> Date: Mon, 20 Apr 2026 11:53:04 +0900 Subject: [PATCH 124/150] [Bugfix] Properly initialize `PerTensorScaleParameter` for fused-on-disk checkpoints (#39765) Signed-off-by: Hemmi Shinichi Signed-off-by: Shinichi Hemmi <50256998+Alnusjaponica@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- vllm/model_executor/layers/linear.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 975fedabd67..b7136097811 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -916,9 +916,15 @@ class MergedColumnParallelLinear(ColumnParallelLinear): loaded_weight=loaded_weight, shard_id=idx ) else: - param.load_merged_column_weight( - loaded_weight=loaded_weight, shard_id=0 - ) + # When weights are already fused on disk (e.g. Phi-3's + # gate_up_proj), there is only a single scale for the + # entire fused matrix. Fill all slots with this scale + # to ensure that any subsequent reduction (like .max()) + # works correctly while preserving the parameter shape. + for idx in range(param.data.shape[0]): + param.load_merged_column_weight( + loaded_weight=loaded_weight, shard_id=idx + ) return elif type(param) in (RowvLLMParameter, BasevLLMParameter): param.load_merged_column_weight(loaded_weight=loaded_weight) @@ -1130,9 +1136,15 @@ class QKVParallelLinear(ColumnParallelLinear): self.validate_shard_id(loaded_shard_id) if loaded_shard_id is None: # special case for certain models if isinstance(param, PerTensorScaleParameter): - param.load_qkv_weight( - loaded_weight=loaded_weight, shard_id=0, tp_rank=self.tp_rank - ) + # When weights are already fused on disk (e.g. Phi-3's + # qkv_proj), there is only a single scale for the entire + # fused matrix. Fill all slots (q, k, v) with this scale + # to ensure that any subsequent reduction (like .max()) + # works correctly while preserving the parameter shape. + for idx in range(param.data.shape[0]): + param.load_qkv_weight( + loaded_weight=loaded_weight, shard_id=idx, tp_rank=self.tp_rank + ) return elif type(param) in (RowvLLMParameter, BasevLLMParameter): param.load_qkv_weight(loaded_weight=loaded_weight, tp_rank=self.tp_rank) From 6e10cb54f69bf26cd4986cda3faa5eb07f25875e Mon Sep 17 00:00:00 2001 From: Hoang Nguyen <118159510+hnt2601@users.noreply.github.com> Date: Mon, 20 Apr 2026 10:24:52 +0700 Subject: [PATCH 125/150] [Bugfix][Responses API] Fix streaming tool calls on /v1/responses (#39892) Signed-off-by: Hoang Nguyen <118159510+hnt2601@users.noreply.github.com> Co-authored-by: Claude --- .../test_gemma4_responses_adjust_request.py | 116 ++++++++++++++++++ vllm/tool_parsers/abstract_tool_parser.py | 21 ++-- vllm/tool_parsers/gemma4_tool_parser.py | 13 +- 3 files changed, 137 insertions(+), 13 deletions(-) create mode 100644 tests/tool_use/test_gemma4_responses_adjust_request.py diff --git a/tests/tool_use/test_gemma4_responses_adjust_request.py b/tests/tool_use/test_gemma4_responses_adjust_request.py new file mode 100644 index 00000000000..e08896ee323 --- /dev/null +++ b/tests/tool_use/test_gemma4_responses_adjust_request.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for Responses API tool-calling request adjustment. + +Covers two bugs on the ``/v1/responses`` path that broke streaming tool +calling for parsers relying on special-token delimiters (Gemma4): + +1. :class:`Gemma4ToolParser.adjust_request` used an + ``isinstance(request, ChatCompletionRequest)`` guard, so a + :class:`ResponsesRequest` with tools never had + ``skip_special_tokens`` flipped to ``False``. The default (``True``) + stripped ``<|tool_call>`` / ```` delimiters, causing + :meth:`Gemma4ToolParser.extract_tool_calls_streaming` to fall through + to the content branch and leak the raw ``call:fn{...}`` body via + ``response.output_text.delta``. + +2. :meth:`ToolParser.adjust_request` built + :class:`ResponseTextConfig` in two steps (bare constructor then + ``.format = ...``). Under Pydantic v2 the later assignment is not + tracked in ``__fields_set__``, which can drop the nested config from + ``model_dump``. It also passed a ``description`` kwarg carrying the + wrong-purpose string ``"Response format for tool calling"``. +""" + +from __future__ import annotations + +from typing import Any + +from openai.types.responses.tool_param import FunctionToolParam + +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.tool_parsers.abstract_tool_parser import ToolParser +from vllm.tool_parsers.gemma4_tool_parser import Gemma4ToolParser + + +def _get_weather_tool() -> FunctionToolParam: + return FunctionToolParam( + type="function", + name="get_weather", + description="Get current weather for a city", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + strict=True, + ) + + +def _build_responses_request(*, tool_choice: str) -> ResponsesRequest: + return ResponsesRequest( + model="gemma4-test", + input=[{"role": "user", "content": "What is the weather in Hanoi?"}], + tools=[_get_weather_tool()], + tool_choice=tool_choice, + stream=True, + max_output_tokens=200, + ) + + +class _StubTokenizer: + """Minimal tokenizer stub to satisfy ``Gemma4ToolParser.__init__``.""" + + def get_vocab(self) -> dict[str, int]: + return {"<|tool_call>": 256_000, "": 256_001, '<|"|>': 52} + + +def test_gemma4_adjust_request_sets_skip_special_tokens_on_responses() -> None: + """``Gemma4ToolParser.adjust_request`` must flip + ``skip_special_tokens=False`` for both ``ChatCompletionRequest`` and + ``ResponsesRequest`` so that ``<|tool_call>`` delimiters reach the + streaming extractor. The previous + ``isinstance(ChatCompletionRequest)`` guard omitted the Responses + path, causing raw ``call:fn{...}`` text to leak via + ``response.output_text.delta``. + """ + parser = Gemma4ToolParser.__new__(Gemma4ToolParser) + parser.model_tokenizer = _StubTokenizer() + + request = _build_responses_request(tool_choice="auto") + assert request.skip_special_tokens is True, ( + "Precondition: ResponsesRequest.skip_special_tokens default is True" + ) + + Gemma4ToolParser.adjust_request(parser, request) + + assert request.skip_special_tokens is False + + +def test_tool_parser_adjust_request_builds_valid_response_text_config() -> None: + """``ToolParser.adjust_request`` must produce a ``ResponseTextConfig`` + whose dumped form contains the JSON schema under the ``schema`` alias + and does not leak the unrelated ``"Response format for tool calling"`` + description string that the previous two-step construction injected. + """ + parser = ToolParser.__new__(ToolParser) + parser.model_tokenizer = None + + request = _build_responses_request(tool_choice="required") + ToolParser.adjust_request(parser, request) + + assert request.text is not None + assert request.text.format is not None + assert request.text.format.type == "json_schema" + + dump: dict[str, Any] = request.text.model_dump(mode="json", by_alias=True) + fmt = dump.get("format") or {} + assert fmt.get("type") == "json_schema" + assert fmt.get("name") == "tool_calling_response" + assert fmt.get("strict") is True + # Nested config must be present under the alias. Two-step Pydantic v2 + # construction could drop it from __fields_set__. + assert "schema" in fmt and isinstance(fmt["schema"], dict) + # The old code passed a wrong-purpose string; valid field should now + # either be absent or None (the openai-python default). + assert fmt.get("description") in (None, "") diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index bc0c090d8bf..75181d8dfac 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/vllm/tool_parsers/abstract_tool_parser.py @@ -103,13 +103,20 @@ class ToolParser: ) request.response_format = None if isinstance(request, ResponsesRequest): - request.text = ResponseTextConfig() - request.text.format = ResponseFormatTextJSONSchemaConfig( - name="tool_calling_response", - schema=json_schema_from_tool, - type="json_schema", - description="Response format for tool calling", - strict=True, + # Single-shot construction so Pydantic v2 tracks `format` + # in __fields_set__ — assigning to `.format` after the bare + # `ResponseTextConfig()` constructor does not, which can + # drop the nested config from `model_dump`. Also drop the + # `description` kwarg: it is not a field on + # ResponseFormatTextJSONSchemaConfig and was being silently + # passed through as extra. + request.text = ResponseTextConfig( + format=ResponseFormatTextJSONSchemaConfig( + type="json_schema", + name="tool_calling_response", + schema=json_schema_from_tool, + strict=True, + ) ) return request diff --git a/vllm/tool_parsers/gemma4_tool_parser.py b/vllm/tool_parsers/gemma4_tool_parser.py index ac48ef26cc1..a5ff2bcf8fc 100644 --- a/vllm/tool_parsers/gemma4_tool_parser.py +++ b/vllm/tool_parsers/gemma4_tool_parser.py @@ -360,12 +360,13 @@ class Gemma4ToolParser(ToolParser): self, request: ChatCompletionRequest | ResponsesRequest ) -> ChatCompletionRequest | ResponsesRequest: request = super().adjust_request(request) - if ( - isinstance(request, ChatCompletionRequest) - and request.tools - and request.tool_choice != "none" - ): - # Don't skip special tokens — <|tool_call> etc. are needed + if request.tools and request.tool_choice != "none": + # Don't skip special tokens — <|tool_call> etc. are needed for + # the parser to detect tool calls. Apply to BOTH + # ChatCompletionRequest and ResponsesRequest (the previous + # isinstance(ChatCompletionRequest) guard caused tool-call + # delimiters to be stripped on /v1/responses, leaking raw + # `call:fn{...}` text via output_text.delta). request.skip_special_tokens = False return request From 67ed01c353f522f3265b95f8a01295f1d99d2969 Mon Sep 17 00:00:00 2001 From: Yuan Tang Date: Mon, 20 Apr 2026 00:17:30 -0400 Subject: [PATCH 126/150] fix: Do not make function calls when request has no tools for /v1/responses (#40314) Signed-off-by: Yuan Tang --- vllm/entrypoints/openai/responses/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index 66239289f73..ad94c8d8dec 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -264,7 +264,7 @@ def convert_tool_responses_to_completions_format(tool: dict) -> dict: def construct_tool_dicts( tools: list[Tool], tool_choice: ToolChoice ) -> list[dict[str, Any]] | None: - if tools is None or (tool_choice == "none"): + if not tools or (tool_choice == "none"): tool_dicts = None else: tool_dicts = [ From 8936118134d0547fa1cc78adab2d03edd6d3dc48 Mon Sep 17 00:00:00 2001 From: Tao He Date: Mon, 20 Apr 2026 12:28:19 +0800 Subject: [PATCH 127/150] [Qwen][Bugfix] Fixes sigmoid activation in torch impl of RMSNormGated. (#40245) Signed-off-by: Tao He --- vllm/model_executor/layers/layernorm.py | 7 +++++-- vllm/model_executor/layers/mamba/gdn_linear_attn.py | 8 ++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 9afc4c9c08d..e4d2d2be090 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -478,9 +478,12 @@ class RMSNormGated(CustomOp): weight = self.weight.float() z = z.float() if z is not None else None + assert self.activation in ["silu", "sigmoid", "swish"] + act_fn = F.sigmoid if self.activation == "sigmoid" else F.silu + # Apply gating before normalization if needed if z is not None and not self.norm_before_gate: - x = x * F.silu(z) + x = x * act_fn(z) # RMS Normalization if self.group_size is None: @@ -499,7 +502,7 @@ class RMSNormGated(CustomOp): # Apply gating after normalization if needed if z is not None and self.norm_before_gate: - out = out * F.silu(z) + out = out * act_fn(z) return out.to(orig_dtype) diff --git a/vllm/model_executor/layers/mamba/gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn_linear_attn.py index 3d875683d26..70a4794ad54 100644 --- a/vllm/model_executor/layers/mamba/gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn_linear_attn.py @@ -357,11 +357,19 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)}) set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) + output_gate_type = getattr(config, "output_gate_type", "silu") + if output_gate_type == "swish": + output_gate_type = "silu" + assert output_gate_type in ["silu", "swish", "sigmoid"], ( + f"unsupported {output_gate_type=}" + ) + self.norm = RMSNormGated( self.head_v_dim, eps=self.layer_norm_epsilon, group_size=None, norm_before_gate=True, + activation=output_gate_type, device=current_platform.current_device(), ) From 4f4713f96e01ebc799125a670c0bca9a2ae7d59d Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Mon, 20 Apr 2026 13:04:39 +0800 Subject: [PATCH 128/150] [XPU] [torch.compile] Skipping CUDA graph memory estimation to avoid startup errors. (#39977) Signed-off-by: chaojun-zhang Co-authored-by: Kunshang Ji --- vllm/v1/worker/gpu_worker.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 5381649f04d..4a73eddcdb3 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -374,10 +374,14 @@ class Worker(WorkerBase): ) # Profile CUDA graph memory if graphs will be captured. - # Skip on ROCm/HIP as graph pool handles and mem_get_info behave + # Skip on ROCm/HIP/XPU as graph pool handles and mem_get_info behave # differently and can produce incorrect/negative estimates. cudagraph_memory_estimate = 0 - if not self.model_config.enforce_eager and not current_platform.is_rocm(): + if ( + not current_platform.is_rocm() + and self.vllm_config.compilation_config.cudagraph_mode + != CUDAGraphMode.NONE + ): cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory() # Use the pre-cudagraph torch peak to avoid double-counting. From 6097afb9bdf9bf59fec70c04b12ab355daf07b23 Mon Sep 17 00:00:00 2001 From: Julien Denize <40604584+juliendenize@users.noreply.github.com> Date: Mon, 20 Apr 2026 07:25:03 +0200 Subject: [PATCH 129/150] [BUGFIX] Fix Pixtral consolidated format vision weight loading (#39916) Signed-off-by: Julien Denize Signed-off-by: juliendenize --- tests/models/fixtures/ministral_3b_chat.json | 1 + .../multimodal/generation/test_pixtral.py | 40 +++++++++++++++++++ vllm/model_executor/models/pixtral.py | 19 +++++++++ 3 files changed, 60 insertions(+) create mode 100644 tests/models/fixtures/ministral_3b_chat.json diff --git a/tests/models/fixtures/ministral_3b_chat.json b/tests/models/fixtures/ministral_3b_chat.json new file mode 100644 index 00000000000..22dd9527adc --- /dev/null +++ b/tests/models/fixtures/ministral_3b_chat.json @@ -0,0 +1 @@ +[[[4380, 3937, 6122, 1261, 7244, 10575, 18970, 41132, 3923, 1408, 1261, 32656, 4691, 1454, 2246, 22131, 15179, 11521, 17277, 1046, 2], "This image shows a black dog sitting attentively on a wooden surface with its gaze directed straight ahead.", [{"4380": {"logprob": -1.0445597171783447, "rank": 1, "decoded_token": "This"}, "1784": {"logprob": -1.5445597171783447, "rank": 2, "decoded_token": "The"}, "1065": {"logprob": -1.7945597171783447, "rank": 3, "decoded_token": "A"}, "1785": {"logprob": -2.2945597171783447, "rank": 4, "decoded_token": "In"}, "4051": {"logprob": -4.357059478759766, "rank": 5, "decoded_token": "An"}}, {"3937": {"logprob": -0.012832445092499256, "rank": 1, "decoded_token": " image"}, "13083": {"logprob": -6.8253326416015625, "rank": 2, "decoded_token": " picture"}, "16649": {"logprob": -6.8878326416015625, "rank": 3, "decoded_token": " photo"}, "1395": {"logprob": -7.2003326416015625, "rank": 4, "decoded_token": " is"}, "7244": {"logprob": -7.5128326416015625, "rank": 5, "decoded_token": " black"}}, {"6122": {"logprob": -0.2629571557044983, "rank": 1, "decoded_token": " shows"}, "51948": {"logprob": -2.2629570960998535, "rank": 2, "decoded_token": " depicts"}, "6971": {"logprob": -2.5129570960998535, "rank": 3, "decoded_token": " features"}, "25981": {"logprob": -3.6379570960998535, "rank": 4, "decoded_token": " displays"}, "89995": {"logprob": -4.7004570960998535, "rank": 5, "decoded_token": " showc"}}, {"1261": {"logprob": -0.00752826826646924, "rank": 1, "decoded_token": " a"}, "1420": {"logprob": -5.132528305053711, "rank": 2, "decoded_token": " an"}, "1278": {"logprob": -6.632528305053711, "rank": 3, "decoded_token": " the"}, "7244": {"logprob": -11.195028305053711, "rank": 4, "decoded_token": " black"}, "44056": {"logprob": -11.757528305053711, "rank": 5, "decoded_token": "\ta"}}, {"7244": {"logprob": -0.19620661437511444, "rank": 1, "decoded_token": " black"}, "8500": {"logprob": -3.446206569671631, "rank": 2, "decoded_token": " dark"}, "4329": {"logprob": -3.446206569671631, "rank": 3, "decoded_token": " large"}, "6231": {"logprob": -3.446206569671631, "rank": 4, "decoded_token": " close"}, "85596": {"logprob": -4.196206569671631, "rank": 5, "decoded_token": " solemn"}}, {"10575": {"logprob": -0.08389955013990402, "rank": 1, "decoded_token": " dog"}, "119075": {"logprob": -3.14639949798584, "rank": 2, "decoded_token": " Labrador"}, "116572": {"logprob": -4.83389949798584, "rank": 3, "decoded_token": " puppy"}, "1044": {"logprob": -5.14639949798584, "rank": 4, "decoded_token": ","}, "9566": {"logprob": -5.58389949798584, "rank": 5, "decoded_token": " medium"}}, {"18970": {"logprob": -1.0850104093551636, "rank": 1, "decoded_token": " sitting"}, "1454": {"logprob": -1.6475104093551636, "rank": 2, "decoded_token": " with"}, "28528": {"logprob": -2.022510528564453, "rank": 3, "decoded_token": " lying"}, "7283": {"logprob": -2.085010528564453, "rank": 4, "decoded_token": " looking"}, "38235": {"logprob": -2.460010528564453, "rank": 5, "decoded_token": " resting"}}, {"41132": {"logprob": -1.2341952323913574, "rank": 1, "decoded_token": " attent"}, "106534": {"logprob": -1.4841952323913574, "rank": 2, "decoded_token": " calmly"}, "1505": {"logprob": -2.2966952323913574, "rank": 3, "decoded_token": " or"}, "17558": {"logprob": -2.3591952323913574, "rank": 4, "decoded_token": " closely"}, "1408": {"logprob": -2.5466952323913574, "rank": 5, "decoded_token": " on"}}, {"3923": {"logprob": -0.005889324937015772, "rank": 1, "decoded_token": "ively"}, "1556": {"logprob": -6.693389415740967, "rank": 2, "decoded_token": "ive"}, "3929": {"logprob": -6.880889415740967, "rank": 3, "decoded_token": "ently"}, "10980": {"logprob": -7.193389415740967, "rank": 4, "decoded_token": "ibly"}, "14194": {"logprob": -7.943389415740967, "rank": 5, "decoded_token": "antly"}}, {"1408": {"logprob": -0.5112090110778809, "rank": 1, "decoded_token": " on"}, "1454": {"logprob": -1.7612090110778809, "rank": 2, "decoded_token": " with"}, "1321": {"logprob": -2.386209011077881, "rank": 3, "decoded_token": " and"}, "3675": {"logprob": -2.761209011077881, "rank": 4, "decoded_token": " against"}, "1505": {"logprob": -4.198709011077881, "rank": 5, "decoded_token": " or"}}, {"1261": {"logprob": -0.189262256026268, "rank": 1, "decoded_token": " a"}, "2549": {"logprob": -2.1892621517181396, "rank": 2, "decoded_token": " what"}, "32656": {"logprob": -3.8767621517181396, "rank": 3, "decoded_token": " wooden"}, "17253": {"logprob": -4.626762390136719, "rank": 4, "decoded_token": " weather"}, "3403": {"logprob": -4.751762390136719, "rank": 5, "decoded_token": " text"}}, {"32656": {"logprob": -0.9549442529678345, "rank": 1, "decoded_token": " wooden"}, "3403": {"logprob": -1.2049442529678345, "rank": 2, "decoded_token": " text"}, "17253": {"logprob": -1.8924442529678345, "rank": 3, "decoded_token": " weather"}, "44130": {"logprob": -2.267444133758545, "rank": 4, "decoded_token": " rust"}, "16673": {"logprob": -4.142444133758545, "rank": 5, "decoded_token": " rough"}}, {"4691": {"logprob": -0.2555857002735138, "rank": 1, "decoded_token": " surface"}, "3403": {"logprob": -2.6930856704711914, "rank": 2, "decoded_token": " text"}, "1615": {"logprob": -2.8805856704711914, "rank": 3, "decoded_token": " pl"}, "1044": {"logprob": -4.193085670471191, "rank": 4, "decoded_token": ","}, "26228": {"logprob": -4.380585670471191, "rank": 5, "decoded_token": " texture"}}, {"1454": {"logprob": -0.5042062997817993, "rank": 1, "decoded_token": " with"}, "1044": {"logprob": -1.6292062997817993, "rank": 2, "decoded_token": ","}, "1046": {"logprob": -2.2542061805725098, "rank": 3, "decoded_token": "."}, "7283": {"logprob": -3.7542061805725098, "rank": 4, "decoded_token": " looking"}, "1338": {"logprob": -4.44170618057251, "rank": 5, "decoded_token": ".\n\n"}}, {"2246": {"logprob": -1.1512703895568848, "rank": 1, "decoded_token": " its"}, "1261": {"logprob": -1.2137703895568848, "rank": 2, "decoded_token": " a"}, "1420": {"logprob": -2.3387703895568848, "rank": 3, "decoded_token": " an"}, "9924": {"logprob": -2.4637703895568848, "rank": 4, "decoded_token": " wide"}, "12593": {"logprob": -3.7762703895568848, "rank": 5, "decoded_token": " slightly"}}, {"22131": {"logprob": -1.0600039958953857, "rank": 1, "decoded_token": " gaze"}, "5731": {"logprob": -1.4350039958953857, "rank": 2, "decoded_token": " eyes"}, "9924": {"logprob": -2.7475039958953857, "rank": 3, "decoded_token": " wide"}, "14781": {"logprob": -2.9350039958953857, "rank": 4, "decoded_token": " focused"}, "3518": {"logprob": -3.1225039958953857, "rank": 5, "decoded_token": " head"}}, {"15179": {"logprob": -1.0535285472869873, "rank": 1, "decoded_token": " directed"}, "9247": {"logprob": -1.5535285472869873, "rank": 2, "decoded_token": " fixed"}, "14781": {"logprob": -2.4285285472869873, "rank": 3, "decoded_token": " focused"}, "7283": {"logprob": -2.6785285472869873, "rank": 4, "decoded_token": " looking"}, "12593": {"logprob": -2.7410285472869873, "rank": 5, "decoded_token": " slightly"}}, {"11521": {"logprob": -1.4163023233413696, "rank": 1, "decoded_token": " straight"}, "40022": {"logprob": -1.6038023233413696, "rank": 2, "decoded_token": " upward"}, "74606": {"logprob": -1.6663023233413696, "rank": 3, "decoded_token": " upwards"}, "12593": {"logprob": -2.16630220413208, "rank": 4, "decoded_token": " slightly"}, "8848": {"logprob": -2.16630220413208, "rank": 5, "decoded_token": " forward"}}, {"17277": {"logprob": -0.20115239918231964, "rank": 1, "decoded_token": " ahead"}, "8848": {"logprob": -2.4511523246765137, "rank": 2, "decoded_token": " forward"}, "1513": {"logprob": -2.9511523246765137, "rank": 3, "decoded_token": " at"}, "1046": {"logprob": -4.576152324676514, "rank": 4, "decoded_token": "."}, "8994": {"logprob": -4.826152324676514, "rank": 5, "decoded_token": " towards"}}, {"1046": {"logprob": -0.1819322109222412, "rank": 1, "decoded_token": "."}, "1338": {"logprob": -2.619432210922241, "rank": 2, "decoded_token": ".\n\n"}, "1321": {"logprob": -3.119432210922241, "rank": 3, "decoded_token": " and"}, "1044": {"logprob": -4.05693244934082, "rank": 4, "decoded_token": ","}, "1626": {"logprob": -5.18193244934082, "rank": 5, "decoded_token": ".\n"}}, {"2": {"logprob": -2.0367233753204346, "rank": 1, "decoded_token": ""}, "35": {"logprob": -4.6617231369018555, "rank": 2, "decoded_token": "[/THINK]"}, "9": {"logprob": -4.9742231369018555, "rank": 3, "decoded_token": "[TOOL_CALLS]"}, "108349": {"logprob": -5.4742231369018555, "rank": 4, "decoded_token": "\u305d\u306e\u4e00\u65b9\u3067"}, "32": {"logprob": -5.8492231369018555, "rank": 5, "decoded_token": "[ARGS]"}}]], [[1784, 2158, 3937, 6122, 1261, 7244, 10575, 18970, 41132, 3923, 1408, 1261, 32656, 4691, 1454, 1261, 26517, 1321, 14781, 4818, 1338, 1784, 2667, 3937, 51948, 1261, 10726, 1290, 3719, 1307, 122203, 1044, 23745, 3591, 13194, 24361, 27469, 1294, 1278, 7786, 1454, 1295, 3506, 11223, 47260, 1408, 1278, 61263, 1046, 2], "The first image shows a black dog sitting attentively on a wooden surface with a calm and focused expression.\n\nThe second image depicts a scenic view of rugged, snow-capped mountain peaks in the distance with lush green vegetation on the slopes.", [{"1784": {"logprob": -0.6781838536262512, "rank": 1, "decoded_token": "The"}, "1049": {"logprob": -1.1781837940216064, "rank": 2, "decoded_token": "1"}, "69957": {"logprob": -4.0531840324401855, "rank": 3, "decoded_token": "Sure"}, "11745": {"logprob": -4.6781840324401855, "rank": 4, "decoded_token": "Here"}, "1785": {"logprob": -5.1781840324401855, "rank": 5, "decoded_token": "In"}}, {"2158": {"logprob": -0.17460788786411285, "rank": 1, "decoded_token": " first"}, "3937": {"logprob": -2.612107992172241, "rank": 2, "decoded_token": " image"}, "8061": {"logprob": -4.299607753753662, "rank": 3, "decoded_token": " images"}, "7244": {"logprob": -5.487107753753662, "rank": 4, "decoded_token": " black"}, "5662": {"logprob": -5.549607753753662, "rank": 5, "decoded_token": " provided"}}, {"3937": {"logprob": -0.014750940725207329, "rank": 1, "decoded_token": " image"}, "13083": {"logprob": -5.889750957489014, "rank": 2, "decoded_token": " picture"}, "2016": {"logprob": -7.264750957489014, "rank": 3, "decoded_token": " set"}, "16649": {"logprob": -7.764750957489014, "rank": 4, "decoded_token": " photo"}, "1877": {"logprob": -7.764750957489014, "rank": 5, "decoded_token": ":\n"}}, {"6122": {"logprob": -0.29441142082214355, "rank": 1, "decoded_token": " shows"}, "51948": {"logprob": -2.2944114208221436, "rank": 2, "decoded_token": " depicts"}, "1877": {"logprob": -3.4194114208221436, "rank": 3, "decoded_token": ":\n"}, "1395": {"logprob": -3.5444114208221436, "rank": 4, "decoded_token": " is"}, "25981": {"logprob": -3.7319114208221436, "rank": 5, "decoded_token": " displays"}}, {"1261": {"logprob": -0.05616755038499832, "rank": 1, "decoded_token": " a"}, "1420": {"logprob": -4.0561676025390625, "rank": 2, "decoded_token": " an"}, "1877": {"logprob": -5.8061676025390625, "rank": 3, "decoded_token": ":\n"}, "1278": {"logprob": -7.1186676025390625, "rank": 4, "decoded_token": " the"}, "20806": {"logprob": -7.6186676025390625, "rank": 5, "decoded_token": "\uff1a\n"}}, {"7244": {"logprob": -0.3175433278083801, "rank": 1, "decoded_token": " black"}, "4329": {"logprob": -2.4425432682037354, "rank": 2, "decoded_token": " large"}, "8500": {"logprob": -3.3175432682037354, "rank": 3, "decoded_token": " dark"}, "6231": {"logprob": -3.6925432682037354, "rank": 4, "decoded_token": " close"}, "85596": {"logprob": -3.7550432682037354, "rank": 5, "decoded_token": " solemn"}}, {"10575": {"logprob": -0.1241316869854927, "rank": 1, "decoded_token": " dog"}, "119075": {"logprob": -2.624131679534912, "rank": 2, "decoded_token": " Labrador"}, "116572": {"logprob": -4.686631679534912, "rank": 3, "decoded_token": " puppy"}, "15812": {"logprob": -5.436631679534912, "rank": 4, "decoded_token": " Lab"}, "1044": {"logprob": -5.561631679534912, "rank": 5, "decoded_token": ","}}, {"18970": {"logprob": -0.9899296164512634, "rank": 1, "decoded_token": " sitting"}, "7283": {"logprob": -1.8649296760559082, "rank": 2, "decoded_token": " looking"}, "28528": {"logprob": -2.114929676055908, "rank": 3, "decoded_token": " lying"}, "1454": {"logprob": -2.364929676055908, "rank": 4, "decoded_token": " with"}, "38235": {"logprob": -2.489929676055908, "rank": 5, "decoded_token": " resting"}}, {"41132": {"logprob": -1.1735738515853882, "rank": 1, "decoded_token": " attent"}, "106534": {"logprob": -1.7360738515853882, "rank": 2, "decoded_token": " calmly"}, "1408": {"logprob": -1.9235738515853882, "rank": 3, "decoded_token": " on"}, "1505": {"logprob": -2.8610739707946777, "rank": 4, "decoded_token": " or"}, "38263": {"logprob": -3.4860739707946777, "rank": 5, "decoded_token": " quietly"}}, {"3923": {"logprob": -0.011853850446641445, "rank": 1, "decoded_token": "ively"}, "1556": {"logprob": -5.6368536949157715, "rank": 2, "decoded_token": "ive"}, "3929": {"logprob": -6.6993536949157715, "rank": 3, "decoded_token": "ently"}, "10980": {"logprob": -6.6993536949157715, "rank": 4, "decoded_token": "ibly"}, "14194": {"logprob": -7.2618536949157715, "rank": 5, "decoded_token": "antly"}}, {"1408": {"logprob": -0.3410560190677643, "rank": 1, "decoded_token": " on"}, "1454": {"logprob": -2.4660561084747314, "rank": 2, "decoded_token": " with"}, "1321": {"logprob": -2.9660561084747314, "rank": 3, "decoded_token": " and"}, "3675": {"logprob": -3.2160561084747314, "rank": 4, "decoded_token": " against"}, "25644": {"logprob": -3.7785561084747314, "rank": 5, "decoded_token": " beside"}}, {"1261": {"logprob": -0.16969695687294006, "rank": 1, "decoded_token": " a"}, "2549": {"logprob": -2.7946970462799072, "rank": 2, "decoded_token": " what"}, "32656": {"logprob": -2.9821970462799072, "rank": 3, "decoded_token": " wooden"}, "3403": {"logprob": -5.107196807861328, "rank": 4, "decoded_token": " text"}, "3977": {"logprob": -5.169696807861328, "rank": 5, "decoded_token": " top"}}, {"32656": {"logprob": -0.49769073724746704, "rank": 1, "decoded_token": " wooden"}, "44130": {"logprob": -1.8726906776428223, "rank": 2, "decoded_token": " rust"}, "3403": {"logprob": -2.2476906776428223, "rank": 3, "decoded_token": " text"}, "17253": {"logprob": -2.9976906776428223, "rank": 4, "decoded_token": " weather"}, "16673": {"logprob": -3.8726906776428223, "rank": 5, "decoded_token": " rough"}}, {"4691": {"logprob": -0.42273157835006714, "rank": 1, "decoded_token": " surface"}, "1615": {"logprob": -1.985231637954712, "rank": 2, "decoded_token": " pl"}, "11237": {"logprob": -3.735231637954712, "rank": 3, "decoded_token": " floor"}, "26808": {"logprob": -3.735231637954712, "rank": 4, "decoded_token": " bench"}, "3403": {"logprob": -3.860231637954712, "rank": 5, "decoded_token": " text"}}, {"1454": {"logprob": -0.8123087286949158, "rank": 1, "decoded_token": " with"}, "1338": {"logprob": -1.5623087882995605, "rank": 2, "decoded_token": ".\n\n"}, "1626": {"logprob": -2.3748087882995605, "rank": 3, "decoded_token": ".\n"}, "7283": {"logprob": -2.8123087882995605, "rank": 4, "decoded_token": " looking"}, "1044": {"logprob": -3.5623087882995605, "rank": 5, "decoded_token": ","}}, {"1261": {"logprob": -1.0109002590179443, "rank": 1, "decoded_token": " a"}, "2246": {"logprob": -1.1984002590179443, "rank": 2, "decoded_token": " its"}, "1420": {"logprob": -2.0109002590179443, "rank": 3, "decoded_token": " an"}, "9924": {"logprob": -3.9484002590179443, "rank": 4, "decoded_token": " wide"}, "14781": {"logprob": -4.698400497436523, "rank": 5, "decoded_token": " focused"}}, {"26517": {"logprob": -2.2192542552948, "rank": 1, "decoded_token": " calm"}, "14781": {"logprob": -2.4067542552948, "rank": 2, "decoded_token": " focused"}, "11304": {"logprob": -2.6567542552948, "rank": 3, "decoded_token": " serious"}, "29691": {"logprob": -2.7192542552948, "rank": 4, "decoded_token": " contempl"}, "12593": {"logprob": -3.1567542552948, "rank": 5, "decoded_token": " slightly"}}, {"1321": {"logprob": -0.8886916637420654, "rank": 1, "decoded_token": " and"}, "4818": {"logprob": -1.0136916637420654, "rank": 2, "decoded_token": " expression"}, "1044": {"logprob": -2.2636916637420654, "rank": 3, "decoded_token": ","}, "22131": {"logprob": -3.1386916637420654, "rank": 4, "decoded_token": " gaze"}, "1311": {"logprob": -3.9511916637420654, "rank": 5, "decoded_token": " de"}}, {"14781": {"logprob": -1.208945870399475, "rank": 1, "decoded_token": " focused"}, "38462": {"logprob": -1.583945870399475, "rank": 2, "decoded_token": " curious"}, "97680": {"logprob": -2.9589457511901855, "rank": 3, "decoded_token": " thoughtful"}, "11304": {"logprob": -2.9589457511901855, "rank": 4, "decoded_token": " serious"}, "29691": {"logprob": -3.3964457511901855, "rank": 5, "decoded_token": " contempl"}}, {"4818": {"logprob": -0.16998574137687683, "rank": 1, "decoded_token": " expression"}, "22131": {"logprob": -2.669985771179199, "rank": 2, "decoded_token": " gaze"}, "2985": {"logprob": -3.232485771179199, "rank": 3, "decoded_token": " look"}, "1311": {"logprob": -3.919985771179199, "rank": 4, "decoded_token": " de"}, "13988": {"logprob": -4.982485771179199, "rank": 5, "decoded_token": " appearance"}}, {"1338": {"logprob": -0.4974508285522461, "rank": 1, "decoded_token": ".\n\n"}, "1626": {"logprob": -1.372450828552246, "rank": 2, "decoded_token": ".\n"}, "1046": {"logprob": -3.497450828552246, "rank": 3, "decoded_token": "."}, "1044": {"logprob": -4.122450828552246, "rank": 4, "decoded_token": ","}, "12560": {"logprob": -4.934950828552246, "rank": 5, "decoded_token": "\u0589\n\n"}}, {"1784": {"logprob": -0.05534925311803818, "rank": 1, "decoded_token": "The"}, "1785": {"logprob": -4.305349349975586, "rank": 2, "decoded_token": "In"}, "6958": {"logprob": -5.617849349975586, "rank": 3, "decoded_token": "There"}, "84593": {"logprob": -5.617849349975586, "rank": 4, "decoded_token": "_The"}, "4393": {"logprob": -6.180349349975586, "rank": 5, "decoded_token": "For"}}, {"2667": {"logprob": -0.035550788044929504, "rank": 1, "decoded_token": " second"}, "3937": {"logprob": -6.160550594329834, "rank": 2, "decoded_token": " image"}, "13023": {"logprob": -6.348050594329834, "rank": 3, "decoded_token": "second"}, "6360": {"logprob": -7.660550594329834, "rank": 4, "decoded_token": " description"}, "2158": {"logprob": -7.785550594329834, "rank": 5, "decoded_token": " first"}}, {"3937": {"logprob": -0.045355767011642456, "rank": 1, "decoded_token": " image"}, "13083": {"logprob": -5.857855796813965, "rank": 2, "decoded_token": " picture"}, "16649": {"logprob": -6.295355796813965, "rank": 3, "decoded_token": " photo"}, "2016": {"logprob": -6.482855796813965, "rank": 4, "decoded_token": " set"}, "5662": {"logprob": -6.857855796813965, "rank": 5, "decoded_token": " provided"}}, {"51948": {"logprob": -0.5656330585479736, "rank": 1, "decoded_token": " depicts"}, "25981": {"logprob": -1.8156330585479736, "rank": 2, "decoded_token": " displays"}, "66583": {"logprob": -2.5031330585479736, "rank": 3, "decoded_token": " captures"}, "6971": {"logprob": -2.7531330585479736, "rank": 4, "decoded_token": " features"}, "1395": {"logprob": -3.3781330585479736, "rank": 5, "decoded_token": " is"}}, {"1261": {"logprob": -0.13412977755069733, "rank": 1, "decoded_token": " a"}, "122203": {"logprob": -2.946629762649536, "rank": 2, "decoded_token": " rugged"}, "1420": {"logprob": -3.821629762649536, "rank": 3, "decoded_token": " an"}, "10726": {"logprob": -4.946630001068115, "rank": 4, "decoded_token": " scen"}, "13770": {"logprob": -5.196630001068115, "rank": 5, "decoded_token": " maj"}}, {"10726": {"logprob": -0.7839541435241699, "rank": 1, "decoded_token": " scen"}, "37849": {"logprob": -2.47145414352417, "rank": 2, "decoded_token": " breat"}, "122203": {"logprob": -2.72145414352417, "rank": 3, "decoded_token": " rugged"}, "23874": {"logprob": -2.72145414352417, "rank": 4, "decoded_token": " pictures"}, "15375": {"logprob": -3.47145414352417, "rank": 5, "decoded_token": " vast"}}, {"1290": {"logprob": -0.000259365770034492, "rank": 1, "decoded_token": "ic"}, "2981": {"logprob": -9.187759399414062, "rank": 2, "decoded_token": "ically"}, "1702": {"logprob": -11.250259399414062, "rank": 3, "decoded_token": "ice"}, "16832": {"logprob": -12.375259399414062, "rank": 4, "decoded_token": "...\n"}, "1685": {"logprob": -12.500259399414062, "rank": 5, "decoded_token": "ical"}}, {"3719": {"logprob": -0.9477202892303467, "rank": 1, "decoded_token": " view"}, "24361": {"logprob": -1.3227202892303467, "rank": 2, "decoded_token": " mountain"}, "127945": {"logprob": -1.9477202892303467, "rank": 3, "decoded_token": " mountainous"}, "1044": {"logprob": -3.0102202892303467, "rank": 4, "decoded_token": ","}, "28035": {"logprob": -3.0727202892303467, "rank": 5, "decoded_token": " landscape"}}, {"1307": {"logprob": -0.030216755345463753, "rank": 1, "decoded_token": " of"}, "1562": {"logprob": -4.217716693878174, "rank": 2, "decoded_token": " from"}, "24018": {"logprob": -5.717716693878174, "rank": 3, "decoded_token": " featuring"}, "11050": {"logprob": -6.217716693878174, "rank": 4, "decoded_token": " showing"}, "2015": {"logprob": -6.280216693878174, "rank": 5, "decoded_token": " up"}}, {"122203": {"logprob": -1.1189241409301758, "rank": 1, "decoded_token": " rugged"}, "1261": {"logprob": -1.6189241409301758, "rank": 2, "decoded_token": " a"}, "127945": {"logprob": -2.681424140930176, "rank": 3, "decoded_token": " mountainous"}, "11223": {"logprob": -2.931424140930176, "rank": 4, "decoded_token": " green"}, "23745": {"logprob": -3.056424140930176, "rank": 5, "decoded_token": " snow"}}, {"1044": {"logprob": -1.033378005027771, "rank": 1, "decoded_token": ","}, "24361": {"logprob": -1.158378005027771, "rank": 2, "decoded_token": " mountain"}, "35463": {"logprob": -1.658378005027771, "rank": 3, "decoded_token": " mountains"}, "127945": {"logprob": -2.2833781242370605, "rank": 4, "decoded_token": " mountainous"}, "1321": {"logprob": -4.9708781242370605, "rank": 5, "decoded_token": " and"}}, {"23745": {"logprob": -0.8830229043960571, "rank": 1, "decoded_token": " snow"}, "127945": {"logprob": -1.8830229043960571, "rank": 2, "decoded_token": " mountainous"}, "11223": {"logprob": -1.8830229043960571, "rank": 3, "decoded_token": " green"}, "1394": {"logprob": -2.5705227851867676, "rank": 4, "decoded_token": " for"}, "95746": {"logprob": -3.5080227851867676, "rank": 5, "decoded_token": " rocky"}}, {"3591": {"logprob": -0.5006332397460938, "rank": 1, "decoded_token": "-c"}, "114525": {"logprob": -1.1256332397460938, "rank": 2, "decoded_token": "-covered"}, "18928": {"logprob": -3.5006332397460938, "rank": 3, "decoded_token": "-top"}, "1099": {"logprob": -4.500633239746094, "rank": 4, "decoded_token": "c"}, "24263": {"logprob": -4.938133239746094, "rank": 5, "decoded_token": "-cl"}}, {"13194": {"logprob": -0.007475971709936857, "rank": 1, "decoded_token": "apped"}, "36649": {"logprob": -6.132475852966309, "rank": 2, "decoded_token": "rowned"}, "13234": {"logprob": -6.694975852966309, "rank": 3, "decoded_token": "aped"}, "10261": {"logprob": -6.819975852966309, "rank": 4, "decoded_token": "rest"}, "4681": {"logprob": -7.757475852966309, "rank": 5, "decoded_token": "ored"}}, {"24361": {"logprob": -0.6446366906166077, "rank": 1, "decoded_token": " mountain"}, "35463": {"logprob": -0.8946366906166077, "rank": 2, "decoded_token": " mountains"}, "127945": {"logprob": -3.394636631011963, "rank": 3, "decoded_token": " mountainous"}, "116555": {"logprob": -4.894636631011963, "rank": 4, "decoded_token": " alpine"}, "1321": {"logprob": -4.894636631011963, "rank": 5, "decoded_token": " and"}}, {"27469": {"logprob": -0.12543931603431702, "rank": 1, "decoded_token": " peaks"}, "26236": {"logprob": -2.625439405441284, "rank": 2, "decoded_token": " ranges"}, "103398": {"logprob": -3.937939405441284, "rank": 3, "decoded_token": " ridges"}, "24765": {"logprob": -4.875439167022705, "rank": 4, "decoded_token": " terrain"}, "84497": {"logprob": -5.187939167022705, "rank": 5, "decoded_token": " landscapes"}}, {"1294": {"logprob": -1.4923679828643799, "rank": 1, "decoded_token": " in"}, "1454": {"logprob": -1.8673679828643799, "rank": 2, "decoded_token": " with"}, "29817": {"logprob": -2.55486798286438, "rank": 3, "decoded_token": " surrounded"}, "26619": {"logprob": -2.61736798286438, "rank": 4, "decoded_token": " rising"}, "1321": {"logprob": -2.61736798286438, "rank": 5, "decoded_token": " and"}}, {"1278": {"logprob": -0.3165818154811859, "rank": 1, "decoded_token": " the"}, "1261": {"logprob": -1.4415818452835083, "rank": 2, "decoded_token": " a"}, "1420": {"logprob": -4.316581726074219, "rank": 3, "decoded_token": " an"}, "5561": {"logprob": -5.754081726074219, "rank": 4, "decoded_token": " various"}, "2549": {"logprob": -6.066581726074219, "rank": 5, "decoded_token": " what"}}, {"7786": {"logprob": -0.5495827198028564, "rank": 1, "decoded_token": " distance"}, "7042": {"logprob": -1.0495827198028564, "rank": 2, "decoded_token": " background"}, "30594": {"logprob": -3.1745827198028564, "rank": 3, "decoded_token": " distant"}, "115381": {"logprob": -5.237082481384277, "rank": 4, "decoded_token": " Alps"}, "92504": {"logprob": -5.237082481384277, "rank": 5, "decoded_token": " backdrop"}}, {"1454": {"logprob": -0.6429675817489624, "rank": 1, "decoded_token": " with"}, "1044": {"logprob": -2.205467700958252, "rank": 2, "decoded_token": ","}, "3675": {"logprob": -2.330467700958252, "rank": 3, "decoded_token": " against"}, "1046": {"logprob": -3.330467700958252, "rank": 4, "decoded_token": "."}, "2136": {"logprob": -3.392967700958252, "rank": 5, "decoded_token": " over"}}, {"1295": {"logprob": -1.4274311065673828, "rank": 1, "decoded_token": " l"}, "1261": {"logprob": -1.6774311065673828, "rank": 2, "decoded_token": " a"}, "11223": {"logprob": -1.9274311065673828, "rank": 3, "decoded_token": " green"}, "47147": {"logprob": -3.239931106567383, "rank": 4, "decoded_token": " steep"}, "50373": {"logprob": -3.302431106567383, "rank": 5, "decoded_token": " patches"}}, {"3506": {"logprob": -0.0006933192489668727, "rank": 1, "decoded_token": "ush"}, "16938": {"logprob": -8.563193321228027, "rank": 2, "decoded_token": "usher"}, "1374": {"logprob": -9.188193321228027, "rank": 3, "decoded_token": "us"}, "90716": {"logprob": -9.563193321228027, "rank": 4, "decoded_token": "USH"}, "5245": {"logprob": -9.688193321228027, "rank": 5, "decoded_token": "acy"}}, {"11223": {"logprob": -0.5139672160148621, "rank": 1, "decoded_token": " green"}, "1044": {"logprob": -1.5764672756195068, "rank": 2, "decoded_token": ","}, "4174": {"logprob": -2.451467275619507, "rank": 3, "decoded_token": " gre"}, "47260": {"logprob": -3.451467275619507, "rank": 4, "decoded_token": " vegetation"}, "1394": {"logprob": -4.763967037200928, "rank": 5, "decoded_token": " for"}}, {"47260": {"logprob": -1.0371246337890625, "rank": 1, "decoded_token": " vegetation"}, "61263": {"logprob": -1.8496246337890625, "rank": 2, "decoded_token": " slopes"}, "94549": {"logprob": -1.9121246337890625, "rank": 3, "decoded_token": " valleys"}, "50373": {"logprob": -3.4121246337890625, "rank": 4, "decoded_token": " patches"}, "4953": {"logprob": -3.5371246337890625, "rank": 5, "decoded_token": " lower"}}, {"1408": {"logprob": -1.3317866325378418, "rank": 1, "decoded_token": " on"}, "1294": {"logprob": -1.7067866325378418, "rank": 2, "decoded_token": " in"}, "1321": {"logprob": -2.394286632537842, "rank": 3, "decoded_token": " and"}, "5956": {"logprob": -2.644286632537842, "rank": 4, "decoded_token": " below"}, "1513": {"logprob": -2.706786632537842, "rank": 5, "decoded_token": " at"}}, {"1278": {"logprob": -0.6638099551200867, "rank": 1, "decoded_token": " the"}, "1261": {"logprob": -2.1638100147247314, "rank": 2, "decoded_token": " a"}, "47147": {"logprob": -2.6013100147247314, "rank": 3, "decoded_token": " steep"}, "4953": {"logprob": -2.9138100147247314, "rank": 4, "decoded_token": " lower"}, "95746": {"logprob": -3.0388100147247314, "rank": 5, "decoded_token": " rocky"}}, {"61263": {"logprob": -1.3242369890213013, "rank": 1, "decoded_token": " slopes"}, "95746": {"logprob": -2.3242368698120117, "rank": 2, "decoded_token": " rocky"}, "4953": {"logprob": -2.3867368698120117, "rank": 3, "decoded_token": " lower"}, "79831": {"logprob": -2.7617368698120117, "rank": 4, "decoded_token": " foreground"}, "47147": {"logprob": -2.8867368698120117, "rank": 5, "decoded_token": " steep"}}, {"1046": {"logprob": -0.7772958278656006, "rank": 1, "decoded_token": "."}, "5956": {"logprob": -1.1522958278656006, "rank": 2, "decoded_token": " below"}, "1294": {"logprob": -2.3397958278656006, "rank": 3, "decoded_token": " in"}, "1321": {"logprob": -3.5897958278656006, "rank": 4, "decoded_token": " and"}, "1307": {"logprob": -4.27729606628418, "rank": 5, "decoded_token": " of"}}, {"2": {"logprob": -0.031993232667446136, "rank": 1, "decoded_token": ""}, "3730": {"logprob": -7.281993389129639, "rank": 2, "decoded_token": " There"}, "1531": {"logprob": -7.281993389129639, "rank": 3, "decoded_token": " The"}, "2409": {"logprob": -8.65699291229248, "rank": 4, "decoded_token": " This"}, "2157": {"logprob": -8.65699291229248, "rank": 5, "decoded_token": " It"}}]], [[1049, 1046, 1531, 2158, 3937, 6122, 1261, 7244, 10575, 18970, 41132, 3923, 1408, 1261, 32656, 4691, 1626, 1256, 1462, 1531, 2667, 3937, 1319, 3715, 4326, 1294, 2143, 4098, 1041, 10249, 1317, 1402, 1261, 24361, 28035, 1454, 122203, 24765, 1321, 30594, 27469, 1338, 1050, 1046, 1531, 5888, 3937, 51948, 1261, 2169, 2509, 29397, 13327, 1454, 26905, 22140, 11981, 1278, 46422, 1321, 1261, 92731, 2965, 19710, 4837, 1278, 46422, 1338, 1051, 1046, 1531, 12432, 3937, 6971, 1261, 53301, 59396, 3549, 121040, 1536, 23170, 1321, 16429, 1294, 1261, 23874, 1872, 41730, 9436, 1338, 1052, 1046, 1531, 19723, 3937, 1319, 3715, 24512, 1435, 1651, 1278, 5719, 4546, 1041, 2168, 1402, 1278, 1925, 1454, 1278, 10575, 2790, 1044, 1809, 4136, 1494, 1681, 5314, 5055, 1044, 3226, 2190, 1261, 2801, 6468, 1693, 1729, 3369, 1278, 3629, 75275, 1877, 1256, 1462, 1531, 5719, 5662, 3937, 1395, 1261, 7244, 10575, 7283, 2015, 1454, 1261, 26517, 4818, 1408, 1261, 44130, 1290, 32656, 1615, 2395, 7042, 1338, 2892, 38695, 1044, 3226, 2190, 1278, 6298, 6360, 1394, 1278, 5662, 7244, 10575, 3937, 1877, 1065, 7244, 10575, 11589, 1264, 1935, 3929, 40022, 1454, 1261, 26517, 4818, 1408, 1261, 44130, 1290, 32656, 1615, 2395, 7042, 1046, 2], "1. The first image shows a black dog sitting attentively on a wooden surface.\n - The second image (not shown in your question) appears to be a mountain landscape with rugged terrain and distant peaks.\n\n2. The third image depicts a serene beach scene with gentle waves meeting the shore and a lone person walking along the shore.\n\n3. The fourth image features a winding gravel path bordered by grass and trees in a picturesque outdoor setting.\n\n4. The fifth image (not applicable as per the initial request) would be the one with the dog again, but since it's already described, here\u2019s a different focus if we consider the following hypothetical:\n - The initial provided image is a black dog looking up with a calm expression on a rustic wooden plank background.\n\nTo clarify, here\u2019s the correct description for the provided black dog image:\nA black dog gazes intently upward with a calm expression on a rustic wooden plank background.", [{"1049": {"logprob": -0.3098178803920746, "rank": 1, "decoded_token": "1"}, "1784": {"logprob": -2.1848177909851074, "rank": 2, "decoded_token": "The"}, "11745": {"logprob": -2.6223177909851074, "rank": 3, "decoded_token": "Here"}, "2757": {"logprob": -6.059817790985107, "rank": 4, "decoded_token": "It"}, "69957": {"logprob": -6.122317790985107, "rank": 5, "decoded_token": "Sure"}}, {"1046": {"logprob": -0.16265615820884705, "rank": 1, "decoded_token": "."}, "1626": {"logprob": -6.475156307220459, "rank": 2, "decoded_token": ".\n"}, "1314": {"logprob": -6.975156307220459, "rank": 3, "decoded_token": "st"}, "1319": {"logprob": -7.412656307220459, "rank": 4, "decoded_token": " ("}, "27": {"logprob": -7.725156307220459, "rank": 5, "decoded_token": ""}}, {"1531": {"logprob": -0.5334714651107788, "rank": 1, "decoded_token": " The"}, "11967": {"logprob": -3.9709715843200684, "rank": 2, "decoded_token": " Image"}, "1349": {"logprob": -3.9709715843200684, "rank": 3, "decoded_token": " A"}, "7610": {"logprob": -3.9709715843200684, "rank": 4, "decoded_token": " First"}, "2048": {"logprob": -4.345971584320068, "rank": 5, "decoded_token": " An"}}, {"2158": {"logprob": -0.4163813889026642, "rank": 1, "decoded_token": " first"}, "3937": {"logprob": -1.2913813591003418, "rank": 2, "decoded_token": " image"}, "7244": {"logprob": -3.853881359100342, "rank": 3, "decoded_token": " black"}, "13083": {"logprob": -4.166381359100342, "rank": 4, "decoded_token": " picture"}, "16649": {"logprob": -5.103881359100342, "rank": 5, "decoded_token": " photo"}}, {"3937": {"logprob": -0.09788376092910767, "rank": 1, "decoded_token": " image"}, "13083": {"logprob": -3.785383701324463, "rank": 2, "decoded_token": " picture"}, "2016": {"logprob": -4.410383701324463, "rank": 3, "decoded_token": " set"}, "1319": {"logprob": -4.472883701324463, "rank": 4, "decoded_token": " ("}, "5662": {"logprob": -4.847883701324463, "rank": 5, "decoded_token": " provided"}}, {"6122": {"logprob": -0.3766213655471802, "rank": 1, "decoded_token": " shows"}, "1395": {"logprob": -2.1266212463378906, "rank": 2, "decoded_token": " is"}, "51948": {"logprob": -2.6891212463378906, "rank": 3, "decoded_token": " depicts"}, "6971": {"logprob": -3.4391212463378906, "rank": 4, "decoded_token": " features"}, "1058": {"logprob": -3.5641212463378906, "rank": 5, "decoded_token": ":"}}, {"1261": {"logprob": -0.04542229697108269, "rank": 1, "decoded_token": " a"}, "1420": {"logprob": -3.170422315597534, "rank": 2, "decoded_token": " an"}, "1278": {"logprob": -6.732922077178955, "rank": 3, "decoded_token": " the"}, "1877": {"logprob": -8.795422554016113, "rank": 4, "decoded_token": ":\n"}, "7244": {"logprob": -9.232922554016113, "rank": 5, "decoded_token": " black"}}, {"7244": {"logprob": -0.45382726192474365, "rank": 1, "decoded_token": " black"}, "4329": {"logprob": -2.266327381134033, "rank": 2, "decoded_token": " large"}, "85596": {"logprob": -3.141327381134033, "rank": 3, "decoded_token": " solemn"}, "8500": {"logprob": -3.203827381134033, "rank": 4, "decoded_token": " dark"}, "16450": {"logprob": -3.766327381134033, "rank": 5, "decoded_token": " sle"}}, {"10575": {"logprob": -0.09157121181488037, "rank": 1, "decoded_token": " dog"}, "119075": {"logprob": -3.27907133102417, "rank": 2, "decoded_token": " Labrador"}, "116572": {"logprob": -4.40407133102417, "rank": 3, "decoded_token": " puppy"}, "15812": {"logprob": -5.40407133102417, "rank": 4, "decoded_token": " Lab"}, "1044": {"logprob": -5.52907133102417, "rank": 5, "decoded_token": ","}}, {"18970": {"logprob": -0.9958467483520508, "rank": 1, "decoded_token": " sitting"}, "7283": {"logprob": -2.058346748352051, "rank": 2, "decoded_token": " looking"}, "28528": {"logprob": -2.120846748352051, "rank": 3, "decoded_token": " lying"}, "38235": {"logprob": -2.183346748352051, "rank": 4, "decoded_token": " resting"}, "1454": {"logprob": -2.558346748352051, "rank": 5, "decoded_token": " with"}}, {"41132": {"logprob": -1.2059022188186646, "rank": 1, "decoded_token": " attent"}, "1408": {"logprob": -1.8309022188186646, "rank": 2, "decoded_token": " on"}, "106534": {"logprob": -1.8934022188186646, "rank": 3, "decoded_token": " calmly"}, "1321": {"logprob": -3.268402099609375, "rank": 4, "decoded_token": " and"}, "1454": {"logprob": -3.330902099609375, "rank": 5, "decoded_token": " with"}}, {"3923": {"logprob": -0.005489394534379244, "rank": 1, "decoded_token": "ively"}, "1556": {"logprob": -5.505489349365234, "rank": 2, "decoded_token": "ive"}, "3929": {"logprob": -7.505489349365234, "rank": 3, "decoded_token": "ently"}, "10980": {"logprob": -8.505489349365234, "rank": 4, "decoded_token": "ibly"}, "14194": {"logprob": -9.130489349365234, "rank": 5, "decoded_token": "antly"}}, {"1408": {"logprob": -0.3085399270057678, "rank": 1, "decoded_token": " on"}, "1454": {"logprob": -2.433539867401123, "rank": 2, "decoded_token": " with"}, "1321": {"logprob": -3.308539867401123, "rank": 3, "decoded_token": " and"}, "25644": {"logprob": -3.496039867401123, "rank": 4, "decoded_token": " beside"}, "3675": {"logprob": -3.746039867401123, "rank": 5, "decoded_token": " against"}}, {"1261": {"logprob": -0.17512622475624084, "rank": 1, "decoded_token": " a"}, "32656": {"logprob": -2.550126314163208, "rank": 2, "decoded_token": " wooden"}, "2549": {"logprob": -3.800126314163208, "rank": 3, "decoded_token": " what"}, "44130": {"logprob": -4.675126075744629, "rank": 4, "decoded_token": " rust"}, "17253": {"logprob": -4.862626075744629, "rank": 5, "decoded_token": " weather"}}, {"32656": {"logprob": -0.5213224291801453, "rank": 1, "decoded_token": " wooden"}, "44130": {"logprob": -1.70882248878479, "rank": 2, "decoded_token": " rust"}, "3403": {"logprob": -2.45882248878479, "rank": 3, "decoded_token": " text"}, "17253": {"logprob": -2.89632248878479, "rank": 4, "decoded_token": " weather"}, "16673": {"logprob": -3.89632248878479, "rank": 5, "decoded_token": " rough"}}, {"4691": {"logprob": -0.7510372996330261, "rank": 1, "decoded_token": " surface"}, "1615": {"logprob": -1.188537359237671, "rank": 2, "decoded_token": " pl"}, "9710": {"logprob": -3.751037359237671, "rank": 3, "decoded_token": " board"}, "3403": {"logprob": -3.751037359237671, "rank": 4, "decoded_token": " text"}, "11237": {"logprob": -3.938537359237671, "rank": 5, "decoded_token": " floor"}}, {"1626": {"logprob": -0.5343737006187439, "rank": 1, "decoded_token": ".\n"}, "1454": {"logprob": -1.4093737602233887, "rank": 2, "decoded_token": " with"}, "1338": {"logprob": -2.7843737602233887, "rank": 3, "decoded_token": ".\n\n"}, "7283": {"logprob": -3.5343737602233887, "rank": 4, "decoded_token": " looking"}, "1044": {"logprob": -3.8468737602233887, "rank": 5, "decoded_token": ","}}, {"1256": {"logprob": -1.0727235078811646, "rank": 1, "decoded_token": " "}, "1050": {"logprob": -1.4477235078811646, "rank": 2, "decoded_token": "2"}, "1293": {"logprob": -2.947723388671875, "rank": 3, "decoded_token": " "}, "6837": {"logprob": -5.260223388671875, "rank": 4, "decoded_token": "\u06f2"}, "1260": {"logprob": -5.822723388671875, "rank": 5, "decoded_token": " "}}, {"1462": {"logprob": -2.970078468322754, "rank": 1, "decoded_token": " -"}, "1319": {"logprob": -3.470078468322754, "rank": 2, "decoded_token": " ("}, "1032": {"logprob": -4.220078468322754, "rank": 3, "decoded_token": " "}, "49958": {"logprob": -4.720078468322754, "rank": 4, "decoded_token": " ```\n"}, "9380": {"logprob": -5.282578468322754, "rank": 5, "decoded_token": " Here"}}, {"1531": {"logprob": -1.9009761810302734, "rank": 1, "decoded_token": " The"}, "10322": {"logprob": -2.9634761810302734, "rank": 2, "decoded_token": " Second"}, "1319": {"logprob": -3.1509761810302734, "rank": 3, "decoded_token": " ("}, "9380": {"logprob": -3.3384761810302734, "rank": 4, "decoded_token": " Here"}, "11967": {"logprob": -3.4634761810302734, "rank": 5, "decoded_token": " Image"}}, {"2667": {"logprob": -0.49421006441116333, "rank": 1, "decoded_token": " second"}, "3937": {"logprob": -2.9942100048065186, "rank": 2, "decoded_token": " image"}, "7244": {"logprob": -3.6192100048065186, "rank": 3, "decoded_token": " black"}, "13827": {"logprob": -3.7442100048065186, "rank": 4, "decoded_token": " subsequent"}, "15115": {"logprob": -3.8692100048065186, "rank": 5, "decoded_token": " detailed"}}, {"3937": {"logprob": -0.938382625579834, "rank": 1, "decoded_token": " image"}, "2016": {"logprob": -2.500882625579834, "rank": 2, "decoded_token": " set"}, "1319": {"logprob": -2.938382625579834, "rank": 3, "decoded_token": " ("}, "5662": {"logprob": -3.000882625579834, "rank": 4, "decoded_token": " provided"}, "7293": {"logprob": -3.438382625579834, "rank": 5, "decoded_token": " actual"}}, {"1319": {"logprob": -2.159437417984009, "rank": 3, "decoded_token": " ("}, "51948": {"logprob": -2.159437417984009, "rank": 1, "decoded_token": " depicts"}, "1395": {"logprob": -2.159437417984009, "rank": 2, "decoded_token": " is"}, "10249": {"logprob": -2.659437417984009, "rank": 4, "decoded_token": " appears"}, "1058": {"logprob": -2.784437417984009, "rank": 5, "decoded_token": ":"}}, {"3715": {"logprob": -2.364915370941162, "rank": 1, "decoded_token": "not"}, "11018": {"logprob": -2.614915370941162, "rank": 2, "decoded_token": "which"}, "2914": {"logprob": -2.927415370941162, "rank": 3, "decoded_token": "inc"}, "39575": {"logprob": -3.239915370941162, "rank": 4, "decoded_token": "mist"}, "3265": {"logprob": -3.302415370941162, "rank": 5, "decoded_token": "the"}}, {"4326": {"logprob": -1.9997965097427368, "rank": 1, "decoded_token": " shown"}, "24512": {"logprob": -2.2497963905334473, "rank": 2, "decoded_token": " applicable"}, "5662": {"logprob": -2.8747963905334473, "rank": 3, "decoded_token": " provided"}, "13874": {"logprob": -2.9372963905334473, "rank": 4, "decoded_token": " visible"}, "1278": {"logprob": -3.1247963905334473, "rank": 5, "decoded_token": " the"}}, {"1294": {"logprob": -1.0730445384979248, "rank": 1, "decoded_token": " in"}, "1041": {"logprob": -2.260544538497925, "rank": 2, "decoded_token": ")"}, "4244": {"logprob": -2.448044538497925, "rank": 3, "decoded_token": "):"}, "3226": {"logprob": -2.635544538497925, "rank": 4, "decoded_token": " here"}, "1435": {"logprob": -3.135544538497925, "rank": 5, "decoded_token": " as"}}, {"2143": {"logprob": -0.9278546571731567, "rank": 1, "decoded_token": " your"}, "1278": {"logprob": -1.2403546571731567, "rank": 2, "decoded_token": " the"}, "4098": {"logprob": -2.177854537963867, "rank": 3, "decoded_token": " question"}, "12705": {"logprob": -3.740354537963867, "rank": 4, "decoded_token": " detail"}, "1593": {"logprob": -3.927854537963867, "rank": 5, "decoded_token": " this"}}, {"4098": {"logprob": -1.8063793182373047, "rank": 1, "decoded_token": " question"}, "4546": {"logprob": -2.3063793182373047, "rank": 2, "decoded_token": " request"}, "4618": {"logprob": -2.4313793182373047, "rank": 3, "decoded_token": " original"}, "7330": {"logprob": -2.5563793182373047, "rank": 4, "decoded_token": " query"}, "5719": {"logprob": -2.8688793182373047, "rank": 5, "decoded_token": " initial"}}, {"1041": {"logprob": -1.411703109741211, "rank": 1, "decoded_token": ")"}, "4244": {"logprob": -2.099203109741211, "rank": 2, "decoded_token": "):"}, "1809": {"logprob": -2.536703109741211, "rank": 3, "decoded_token": " but"}, "1044": {"logprob": -2.661703109741211, "rank": 4, "decoded_token": ","}, "1681": {"logprob": -3.474203109741211, "rank": 5, "decoded_token": "'s"}}, {"10249": {"logprob": -1.607055425643921, "rank": 1, "decoded_token": " appears"}, "51948": {"logprob": -1.982055425643921, "rank": 2, "decoded_token": " depicts"}, "7444": {"logprob": -2.044555425643921, "rank": 3, "decoded_token": " seems"}, "1395": {"logprob": -2.169555425643921, "rank": 4, "decoded_token": " is"}, "2168": {"logprob": -2.982055425643921, "rank": 5, "decoded_token": " would"}}, {"1317": {"logprob": -0.1792934685945511, "rank": 1, "decoded_token": " to"}, "51723": {"logprob": -4.116793632507324, "rank": 2, "decoded_token": " unrelated"}, "5711": {"logprob": -4.616793632507324, "rank": 3, "decoded_token": " mis"}, "73751": {"logprob": -4.804293632507324, "rank": 4, "decoded_token": " incorrectly"}, "1605": {"logprob": -4.866793632507324, "rank": 5, "decoded_token": " not"}}, {"1402": {"logprob": -0.8325714468955994, "rank": 1, "decoded_token": " be"}, "17767": {"logprob": -1.4575715065002441, "rank": 2, "decoded_token": " depict"}, "7326": {"logprob": -3.395071506500244, "rank": 3, "decoded_token": " feature"}, "1736": {"logprob": -3.520071506500244, "rank": 4, "decoded_token": " have"}, "32688": {"logprob": -3.582571506500244, "rank": 5, "decoded_token": " illustrate"}}, {"1261": {"logprob": -1.2619296312332153, "rank": 1, "decoded_token": " a"}, "1420": {"logprob": -2.824429512023926, "rank": 2, "decoded_token": " an"}, "51723": {"logprob": -3.261929512023926, "rank": 3, "decoded_token": " unrelated"}, "25342": {"logprob": -3.511929512023926, "rank": 4, "decoded_token": " incorrect"}, "1307": {"logprob": -3.511929512023926, "rank": 5, "decoded_token": " of"}}, {"24361": {"logprob": -2.2388875484466553, "rank": 1, "decoded_token": " mountain"}, "10726": {"logprob": -2.5513875484466553, "rank": 2, "decoded_token": " scen"}, "127945": {"logprob": -3.3013875484466553, "rank": 3, "decoded_token": " mountainous"}, "15115": {"logprob": -3.6138875484466553, "rank": 4, "decoded_token": " detailed"}, "3719": {"logprob": -3.7388875484466553, "rank": 5, "decoded_token": " view"}}, {"28035": {"logprob": -0.8160803318023682, "rank": 1, "decoded_token": " landscape"}, "4521": {"logprob": -1.2535803318023682, "rank": 2, "decoded_token": " range"}, "3719": {"logprob": -2.691080331802368, "rank": 3, "decoded_token": " view"}, "13327": {"logprob": -3.253580331802368, "rank": 4, "decoded_token": " scene"}, "1454": {"logprob": -3.566080331802368, "rank": 5, "decoded_token": " with"}}, {"1454": {"logprob": -0.6144524812698364, "rank": 1, "decoded_token": " with"}, "1626": {"logprob": -2.051952362060547, "rank": 2, "decoded_token": ".\n"}, "1338": {"logprob": -2.176952362060547, "rank": 3, "decoded_token": ".\n\n"}, "3719": {"logprob": -3.926952362060547, "rank": 4, "decoded_token": " view"}, "1562": {"logprob": -3.989452362060547, "rank": 5, "decoded_token": " from"}}, {"122203": {"logprob": -1.7088322639465332, "rank": 1, "decoded_token": " rugged"}, "23745": {"logprob": -2.208832263946533, "rank": 2, "decoded_token": " snow"}, "11223": {"logprob": -2.521332263946533, "rank": 3, "decoded_token": " green"}, "27469": {"logprob": -2.833832263946533, "rank": 4, "decoded_token": " peaks"}, "47147": {"logprob": -2.896332263946533, "rank": 5, "decoded_token": " steep"}}, {"24765": {"logprob": -0.6277848482131958, "rank": 1, "decoded_token": " terrain"}, "27469": {"logprob": -1.2527848482131958, "rank": 2, "decoded_token": " peaks"}, "1044": {"logprob": -2.9402847290039062, "rank": 3, "decoded_token": ","}, "130655": {"logprob": -3.6277847290039062, "rank": 4, "decoded_token": " cliffs"}, "57912": {"logprob": -4.002784729003906, "rank": 5, "decoded_token": " terrains"}}, {"1321": {"logprob": -0.6374254822731018, "rank": 1, "decoded_token": " and"}, "1626": {"logprob": -1.762425422668457, "rank": 2, "decoded_token": ".\n"}, "1338": {"logprob": -1.762425422668457, "rank": 3, "decoded_token": ".\n\n"}, "1294": {"logprob": -3.887425422668457, "rank": 4, "decoded_token": " in"}, "2425": {"logprob": -4.012425422668457, "rank": 5, "decoded_token": " under"}}, {"30594": {"logprob": -1.8202354907989502, "rank": 1, "decoded_token": " distant"}, "23745": {"logprob": -2.13273549079895, "rank": 2, "decoded_token": " snow"}, "27469": {"logprob": -2.25773549079895, "rank": 3, "decoded_token": " peaks"}, "11223": {"logprob": -2.94523549079895, "rank": 4, "decoded_token": " green"}, "11692": {"logprob": -3.07023549079895, "rank": 5, "decoded_token": " mist"}}, {"27469": {"logprob": -0.40320226550102234, "rank": 1, "decoded_token": " peaks"}, "23745": {"logprob": -2.2157022953033447, "rank": 2, "decoded_token": " snow"}, "35463": {"logprob": -3.4032022953033447, "rank": 3, "decoded_token": " mountains"}, "24361": {"logprob": -3.4032022953033447, "rank": 4, "decoded_token": " mountain"}, "46866": {"logprob": -4.090702056884766, "rank": 5, "decoded_token": " clouds"}}, {"1338": {"logprob": -0.5224027633666992, "rank": 1, "decoded_token": ".\n\n"}, "1626": {"logprob": -1.3349027633666992, "rank": 2, "decoded_token": ".\n"}, "2425": {"logprob": -3.397402763366699, "rank": 3, "decoded_token": " under"}, "1294": {"logprob": -4.084902763366699, "rank": 4, "decoded_token": " in"}, "13875": {"logprob": -4.334902763366699, "rank": 5, "decoded_token": " covered"}}, {"1050": {"logprob": -0.6029570698738098, "rank": 1, "decoded_token": "2"}, "1256": {"logprob": -1.352957010269165, "rank": 2, "decoded_token": " "}, "11745": {"logprob": -3.227957010269165, "rank": 3, "decoded_token": "Here"}, "1293": {"logprob": -3.915457010269165, "rank": 4, "decoded_token": " "}, "4393": {"logprob": -4.602957248687744, "rank": 5, "decoded_token": "For"}}, {"1046": {"logprob": -0.19250261783599854, "rank": 1, "decoded_token": "."}, "1319": {"logprob": -3.567502498626709, "rank": 2, "decoded_token": " ("}, "1626": {"logprob": -4.505002498626709, "rank": 3, "decoded_token": ".\n"}, "1877": {"logprob": -5.942502498626709, "rank": 4, "decoded_token": ":\n"}, "26667": {"logprob": -6.442502498626709, "rank": 5, "decoded_token": ".("}}, {"1531": {"logprob": -1.897325873374939, "rank": 1, "decoded_token": " The"}, "9380": {"logprob": -1.897325873374939, "rank": 2, "decoded_token": " Here"}, "1319": {"logprob": -2.0848259925842285, "rank": 3, "decoded_token": " ("}, "2898": {"logprob": -2.8973259925842285, "rank": 4, "decoded_token": " For"}, "10322": {"logprob": -3.3348259925842285, "rank": 5, "decoded_token": " Second"}}, {"5888": {"logprob": -1.483720064163208, "rank": 1, "decoded_token": " third"}, "2667": {"logprob": -2.296220064163208, "rank": 2, "decoded_token": " second"}, "5662": {"logprob": -2.858720064163208, "rank": 3, "decoded_token": " provided"}, "3937": {"logprob": -3.296220064163208, "rank": 4, "decoded_token": " image"}, "6298": {"logprob": -3.421220064163208, "rank": 5, "decoded_token": " correct"}}, {"3937": {"logprob": -0.6041796207427979, "rank": 1, "decoded_token": " image"}, "1319": {"logprob": -2.604179620742798, "rank": 2, "decoded_token": " ("}, "2016": {"logprob": -2.729179620742798, "rank": 3, "decoded_token": " set"}, "1321": {"logprob": -3.916679620742798, "rank": 4, "decoded_token": " and"}, "13083": {"logprob": -3.979179620742798, "rank": 5, "decoded_token": " picture"}}, {"51948": {"logprob": -1.6496541500091553, "rank": 1, "decoded_token": " depicts"}, "1319": {"logprob": -1.8371541500091553, "rank": 2, "decoded_token": " ("}, "1058": {"logprob": -1.9621541500091553, "rank": 3, "decoded_token": ":"}, "25981": {"logprob": -2.4621541500091553, "rank": 4, "decoded_token": " displays"}, "1395": {"logprob": -2.7746541500091553, "rank": 5, "decoded_token": " is"}}, {"1261": {"logprob": -0.5954864025115967, "rank": 1, "decoded_token": " a"}, "1420": {"logprob": -2.5954864025115967, "rank": 2, "decoded_token": " an"}, "22140": {"logprob": -2.6579864025115967, "rank": 3, "decoded_token": " waves"}, "1278": {"logprob": -2.9704864025115967, "rank": 4, "decoded_token": " the"}, "26517": {"logprob": -3.9079864025115967, "rank": 5, "decoded_token": " calm"}}, {"2169": {"logprob": -2.0800817012786865, "rank": 1, "decoded_token": " ser"}, "10726": {"logprob": -2.2050817012786865, "rank": 2, "decoded_token": " scen"}, "29397": {"logprob": -2.3925817012786865, "rank": 3, "decoded_token": " beach"}, "2965": {"logprob": -2.4550817012786865, "rank": 4, "decoded_token": " person"}, "1958": {"logprob": -2.8925817012786865, "rank": 5, "decoded_token": " sur"}}, {"2509": {"logprob": -0.0010151476599276066, "rank": 1, "decoded_token": "ene"}, "1391": {"logprob": -8.188514709472656, "rank": 2, "decoded_token": "if"}, "25863": {"logprob": -8.938514709472656, "rank": 3, "decoded_token": "rated"}, "3414": {"logprob": -9.438514709472656, "rank": 4, "decoded_token": "ena"}, "10049": {"logprob": -9.688514709472656, "rank": 5, "decoded_token": "\u00e8ne"}}, {"29397": {"logprob": -0.5060035586357117, "rank": 1, "decoded_token": " beach"}, "62557": {"logprob": -2.5685036182403564, "rank": 2, "decoded_token": " seas"}, "27208": {"logprob": -2.9435036182403564, "rank": 3, "decoded_token": " ocean"}, "38167": {"logprob": -3.0685036182403564, "rank": 4, "decoded_token": " coastal"}, "13327": {"logprob": -3.2560036182403564, "rank": 5, "decoded_token": " scene"}}, {"13327": {"logprob": -0.8083749413490295, "rank": 1, "decoded_token": " scene"}, "1454": {"logprob": -1.1208748817443848, "rank": 2, "decoded_token": " with"}, "1513": {"logprob": -2.9333748817443848, "rank": 3, "decoded_token": " at"}, "3184": {"logprob": -3.4958748817443848, "rank": 4, "decoded_token": " during"}, "2478": {"logprob": -4.058374881744385, "rank": 5, "decoded_token": " where"}}, {"1454": {"logprob": -0.322070449590683, "rank": 1, "decoded_token": " with"}, "3184": {"logprob": -2.447070360183716, "rank": 2, "decoded_token": " during"}, "1513": {"logprob": -2.759570360183716, "rank": 3, "decoded_token": " at"}, "2425": {"logprob": -3.509570360183716, "rank": 4, "decoded_token": " under"}, "2478": {"logprob": -3.697070360183716, "rank": 5, "decoded_token": " where"}}, {"26905": {"logprob": -1.1276788711547852, "rank": 1, "decoded_token": " gentle"}, "22140": {"logprob": -1.4401788711547852, "rank": 2, "decoded_token": " waves"}, "1261": {"logprob": -2.377678871154785, "rank": 3, "decoded_token": " a"}, "3306": {"logprob": -2.815178871154785, "rank": 4, "decoded_token": " people"}, "3709": {"logprob": -3.252678871154785, "rank": 5, "decoded_token": " small"}}, {"22140": {"logprob": -0.048392221331596375, "rank": 1, "decoded_token": " waves"}, "27208": {"logprob": -3.6733922958374023, "rank": 2, "decoded_token": " ocean"}, "29661": {"logprob": -5.735892295837402, "rank": 3, "decoded_token": " surf"}, "32911": {"logprob": -5.735892295837402, "rank": 4, "decoded_token": " rolling"}, "11196": {"logprob": -6.173392295837402, "rank": 5, "decoded_token": " sea"}}, {"11981": {"logprob": -1.4009065628051758, "rank": 1, "decoded_token": " meeting"}, "1321": {"logprob": -1.5884065628051758, "rank": 2, "decoded_token": " and"}, "86928": {"logprob": -2.088406562805176, "rank": 3, "decoded_token": " crashing"}, "1427": {"logprob": -2.650906562805176, "rank": 4, "decoded_token": " la"}, "44278": {"logprob": -2.838406562805176, "rank": 5, "decoded_token": " approaching"}}, {"1278": {"logprob": -0.8703328967094421, "rank": 1, "decoded_token": " the"}, "100991": {"logprob": -1.370332956314087, "rank": 2, "decoded_token": " sandy"}, "1261": {"logprob": -1.932832956314087, "rank": 3, "decoded_token": " a"}, "14693": {"logprob": -2.932832956314087, "rank": 4, "decoded_token": " sand"}, "46422": {"logprob": -3.557832956314087, "rank": 5, "decoded_token": " shore"}}, {"46422": {"logprob": -0.19113700091838837, "rank": 1, "decoded_token": " shore"}, "100991": {"logprob": -2.7536370754241943, "rank": 2, "decoded_token": " sandy"}, "1627": {"logprob": -3.0036370754241943, "rank": 3, "decoded_token": " sh"}, "14693": {"logprob": -3.3786370754241943, "rank": 4, "decoded_token": " sand"}, "124562": {"logprob": -5.378636837005615, "rank": 5, "decoded_token": " coastline"}}, {"1321": {"logprob": -1.191644549369812, "rank": 1, "decoded_token": " and"}, "3184": {"logprob": -2.0666446685791016, "rank": 2, "decoded_token": " during"}, "1338": {"logprob": -2.3166446685791016, "rank": 3, "decoded_token": ".\n\n"}, "1044": {"logprob": -2.3791446685791016, "rank": 4, "decoded_token": ","}, "1513": {"logprob": -2.5666446685791016, "rank": 5, "decoded_token": " at"}}, {"1261": {"logprob": -0.937598705291748, "rank": 1, "decoded_token": " a"}, "3306": {"logprob": -2.000098705291748, "rank": 2, "decoded_token": " people"}, "29661": {"logprob": -2.625098705291748, "rank": 3, "decoded_token": " surf"}, "2269": {"logprob": -2.937598705291748, "rank": 4, "decoded_token": " some"}, "30594": {"logprob": -3.062598705291748, "rank": 5, "decoded_token": " distant"}}, {"92731": {"logprob": -1.4627695083618164, "rank": 1, "decoded_token": " lone"}, "2965": {"logprob": -1.6502695083618164, "rank": 2, "decoded_token": " person"}, "4517": {"logprob": -2.0877695083618164, "rank": 3, "decoded_token": " few"}, "81249": {"logprob": -2.3377695083618164, "rank": 4, "decoded_token": " silhouette"}, "79013": {"logprob": -2.4627695083618164, "rank": 5, "decoded_token": " solitary"}}, {"2965": {"logprob": -0.9964981079101562, "rank": 1, "decoded_token": " person"}, "8240": {"logprob": -1.1839981079101562, "rank": 2, "decoded_token": " figure"}, "1958": {"logprob": -1.9964981079101562, "rank": 3, "decoded_token": " sur"}, "81249": {"logprob": -3.2464981079101562, "rank": 4, "decoded_token": " silhouette"}, "4597": {"logprob": -3.3089981079101562, "rank": 5, "decoded_token": " individual"}}, {"19710": {"logprob": -1.8438996076583862, "rank": 1, "decoded_token": " walking"}, "15866": {"logprob": -2.093899726867676, "rank": 2, "decoded_token": " standing"}, "6117": {"logprob": -2.093899726867676, "rank": 3, "decoded_token": " near"}, "1285": {"logprob": -2.468899726867676, "rank": 4, "decoded_token": " w"}, "1294": {"logprob": -2.843899726867676, "rank": 5, "decoded_token": " in"}}, {"4837": {"logprob": -1.405532717704773, "rank": 1, "decoded_token": " along"}, "6117": {"logprob": -1.718032717704773, "rank": 2, "decoded_token": " near"}, "1408": {"logprob": -2.5930328369140625, "rank": 3, "decoded_token": " on"}, "8994": {"logprob": -2.7180328369140625, "rank": 4, "decoded_token": " towards"}, "1338": {"logprob": -2.8430328369140625, "rank": 5, "decoded_token": ".\n\n"}}, {"1278": {"logprob": -0.1722739040851593, "rank": 1, "decoded_token": " the"}, "1494": {"logprob": -2.734773874282837, "rank": 2, "decoded_token": " it"}, "1338": {"logprob": -3.484773874282837, "rank": 3, "decoded_token": ".\n\n"}, "1261": {"logprob": -3.984773874282837, "rank": 4, "decoded_token": " a"}, "1626": {"logprob": -4.984774112701416, "rank": 5, "decoded_token": ".\n"}}, {"46422": {"logprob": -1.466295599937439, "rank": 1, "decoded_token": " shore"}, "1627": {"logprob": -1.528795599937439, "rank": 2, "decoded_token": " sh"}, "4180": {"logprob": -1.841295599937439, "rank": 3, "decoded_token": " water"}, "10314": {"logprob": -2.0912957191467285, "rank": 4, "decoded_token": " edge"}, "14693": {"logprob": -3.0912957191467285, "rank": 5, "decoded_token": " sand"}}, {"1338": {"logprob": -0.9661335945129395, "rank": 1, "decoded_token": ".\n\n"}, "2839": {"logprob": -2.0911335945129395, "rank": 2, "decoded_token": "line"}, "1626": {"logprob": -2.1536335945129395, "rank": 3, "decoded_token": ".\n"}, "1513": {"logprob": -2.1536335945129395, "rank": 4, "decoded_token": " at"}, "3184": {"logprob": -2.6536335945129395, "rank": 5, "decoded_token": " during"}}, {"1051": {"logprob": -0.8673074841499329, "rank": 1, "decoded_token": "3"}, "1256": {"logprob": -2.054807424545288, "rank": 2, "decoded_token": " "}, "11745": {"logprob": -2.242307424545288, "rank": 3, "decoded_token": "Here"}, "1052": {"logprob": -2.242307424545288, "rank": 4, "decoded_token": "4"}, "1050": {"logprob": -3.367307424545288, "rank": 5, "decoded_token": "2"}}, {"1046": {"logprob": -0.09165985882282257, "rank": 1, "decoded_token": "."}, "1626": {"logprob": -4.341660022735596, "rank": 2, "decoded_token": ".\n"}, "1319": {"logprob": -4.841660022735596, "rank": 3, "decoded_token": " ("}, "1045": {"logprob": -6.029160022735596, "rank": 4, "decoded_token": "-"}, "1321": {"logprob": -6.529160022735596, "rank": 5, "decoded_token": " and"}}, {"1531": {"logprob": -0.5217734575271606, "rank": 1, "decoded_token": " The"}, "9380": {"logprob": -2.396773338317871, "rank": 2, "decoded_token": " Here"}, "2898": {"logprob": -3.521773338317871, "rank": 3, "decoded_token": " For"}, "45663": {"logprob": -3.959273338317871, "rank": 4, "decoded_token": " Fourth"}, "1319": {"logprob": -4.084273338317871, "rank": 5, "decoded_token": " ("}}, {"12432": {"logprob": -0.41262897849082947, "rank": 1, "decoded_token": " fourth"}, "5662": {"logprob": -3.8501288890838623, "rank": 2, "decoded_token": " provided"}, "5888": {"logprob": -4.037629127502441, "rank": 3, "decoded_token": " third"}, "3937": {"logprob": -4.037629127502441, "rank": 4, "decoded_token": " image"}, "2667": {"logprob": -4.162629127502441, "rank": 5, "decoded_token": " second"}}, {"3937": {"logprob": -0.5903608798980713, "rank": 1, "decoded_token": " image"}, "1319": {"logprob": -2.3403608798980713, "rank": 2, "decoded_token": " ("}, "2016": {"logprob": -2.7153608798980713, "rank": 3, "decoded_token": " set"}, "1321": {"logprob": -3.4653608798980713, "rank": 4, "decoded_token": " and"}, "1925": {"logprob": -4.090360641479492, "rank": 5, "decoded_token": " one"}}, {"6971": {"logprob": -1.8004755973815918, "rank": 1, "decoded_token": " features"}, "1319": {"logprob": -2.175475597381592, "rank": 2, "decoded_token": " ("}, "6122": {"logprob": -2.300475597381592, "rank": 3, "decoded_token": " shows"}, "1395": {"logprob": -2.487975597381592, "rank": 4, "decoded_token": " is"}, "1058": {"logprob": -2.487975597381592, "rank": 5, "decoded_token": ":"}}, {"1261": {"logprob": -0.16380631923675537, "rank": 1, "decoded_token": " a"}, "1420": {"logprob": -2.976306438446045, "rank": 2, "decoded_token": " an"}, "2295": {"logprob": -3.351306438446045, "rank": 3, "decoded_token": " two"}, "1278": {"logprob": -4.413806438446045, "rank": 4, "decoded_token": " the"}, "16429": {"logprob": -4.851306438446045, "rank": 5, "decoded_token": " trees"}}, {"53301": {"logprob": -1.6735851764678955, "rank": 1, "decoded_token": " winding"}, "59396": {"logprob": -1.9235851764678955, "rank": 2, "decoded_token": " gravel"}, "101727": {"logprob": -2.1110851764678955, "rank": 3, "decoded_token": " paved"}, "47945": {"logprob": -3.1110851764678955, "rank": 4, "decoded_token": " dirt"}, "23874": {"logprob": -3.1110851764678955, "rank": 5, "decoded_token": " pictures"}}, {"59396": {"logprob": -0.6121145486831665, "rank": 1, "decoded_token": " gravel"}, "47945": {"logprob": -1.7996145486831665, "rank": 2, "decoded_token": " dirt"}, "101727": {"logprob": -2.612114429473877, "rank": 3, "decoded_token": " paved"}, "3549": {"logprob": -2.799614429473877, "rank": 4, "decoded_token": " path"}, "1044": {"logprob": -2.862114429473877, "rank": 5, "decoded_token": ","}}, {"3549": {"logprob": -0.4002755880355835, "rank": 1, "decoded_token": " path"}, "14801": {"logprob": -1.7752755880355835, "rank": 2, "decoded_token": " pathway"}, "1505": {"logprob": -2.275275707244873, "rank": 3, "decoded_token": " or"}, "33659": {"logprob": -3.962775707244873, "rank": 4, "decoded_token": " trail"}, "9480": {"logprob": -4.337775707244873, "rank": 5, "decoded_token": " road"}}, {"121040": {"logprob": -1.1325652599334717, "rank": 1, "decoded_token": " bordered"}, "121313": {"logprob": -1.6950652599334717, "rank": 2, "decoded_token": " flanked"}, "29817": {"logprob": -2.1950652599334717, "rank": 3, "decoded_token": " surrounded"}, "8924": {"logprob": -2.6950652599334717, "rank": 4, "decoded_token": " leading"}, "1294": {"logprob": -2.8200652599334717, "rank": 5, "decoded_token": " in"}}, {"1536": {"logprob": -0.02329995296895504, "rank": 1, "decoded_token": " by"}, "1454": {"logprob": -4.5233001708984375, "rank": 2, "decoded_token": " with"}, "1408": {"logprob": -4.5858001708984375, "rank": 3, "decoded_token": " on"}, "98245": {"logprob": -8.460800170898438, "rank": 4, "decoded_token": " beautifully"}, "100910": {"logprob": -8.648300170898438, "rank": 5, "decoded_token": " neatly"}}, {"23170": {"logprob": -1.2441778182983398, "rank": 1, "decoded_token": " grass"}, "1295": {"logprob": -1.4941778182983398, "rank": 2, "decoded_token": " l"}, "11223": {"logprob": -1.5566778182983398, "rank": 3, "decoded_token": " green"}, "1261": {"logprob": -3.05667781829834, "rank": 4, "decoded_token": " a"}, "4174": {"logprob": -3.61917781829834, "rank": 5, "decoded_token": " gre"}}, {"1321": {"logprob": -0.8162240982055664, "rank": 1, "decoded_token": " and"}, "1121": {"logprob": -1.4412240982055664, "rank": 2, "decoded_token": "y"}, "1044": {"logprob": -1.8787240982055664, "rank": 3, "decoded_token": ","}, "8924": {"logprob": -3.3162240982055664, "rank": 4, "decoded_token": " leading"}, "1454": {"logprob": -3.5037240982055664, "rank": 5, "decoded_token": " with"}}, {"16429": {"logprob": -1.8316444158554077, "rank": 1, "decoded_token": " trees"}, "87833": {"logprob": -2.4566445350646973, "rank": 2, "decoded_token": " flowering"}, "17744": {"logprob": -2.5816445350646973, "rank": 3, "decoded_token": " blo"}, "1261": {"logprob": -2.6441445350646973, "rank": 4, "decoded_token": " a"}, "29817": {"logprob": -3.0191445350646973, "rank": 5, "decoded_token": " surrounded"}}, {"1294": {"logprob": -1.625852108001709, "rank": 1, "decoded_token": " in"}, "1044": {"logprob": -1.750852108001709, "rank": 2, "decoded_token": ","}, "2425": {"logprob": -2.125852108001709, "rank": 3, "decoded_token": " under"}, "8924": {"logprob": -2.250852108001709, "rank": 4, "decoded_token": " leading"}, "1408": {"logprob": -2.625852108001709, "rank": 5, "decoded_token": " on"}}, {"1261": {"logprob": -0.3194517195224762, "rank": 1, "decoded_token": " a"}, "1278": {"logprob": -2.3194518089294434, "rank": 2, "decoded_token": " the"}, "1420": {"logprob": -2.5694518089294434, "rank": 3, "decoded_token": " an"}, "2549": {"logprob": -2.6319518089294434, "rank": 4, "decoded_token": " what"}, "5346": {"logprob": -5.569451808929443, "rank": 5, "decoded_token": " front"}}, {"23874": {"logprob": -1.2531365156173706, "rank": 1, "decoded_token": " pictures"}, "10726": {"logprob": -1.6906365156173706, "rank": 2, "decoded_token": " scen"}, "54742": {"logprob": -2.12813663482666, "rank": 3, "decoded_token": " peaceful"}, "12097": {"logprob": -2.69063663482666, "rank": 4, "decoded_token": " park"}, "2169": {"logprob": -3.00313663482666, "rank": 5, "decoded_token": " ser"}}, {"1872": {"logprob": -0.0001784403866622597, "rank": 1, "decoded_token": "que"}, "1348": {"logprob": -8.875178337097168, "rank": 2, "decoded_token": "qu"}, "21451": {"logprob": -11.937678337097168, "rank": 3, "decoded_token": "QUE"}, "18954": {"logprob": -12.687678337097168, "rank": 4, "decoded_token": "qu\u00e9"}, "14016": {"logprob": -12.812678337097168, "rank": 5, "decoded_token": "quare"}}, {"41730": {"logprob": -1.2438380718231201, "rank": 1, "decoded_token": " outdoor"}, "12097": {"logprob": -1.4938380718231201, "rank": 2, "decoded_token": " park"}, "1044": {"logprob": -1.8688380718231201, "rank": 3, "decoded_token": ","}, "28035": {"logprob": -2.55633807182312, "rank": 4, "decoded_token": " landscape"}, "6967": {"logprob": -2.68133807182312, "rank": 5, "decoded_token": " natural"}}, {"9436": {"logprob": -0.28005897998809814, "rank": 1, "decoded_token": " setting"}, "6093": {"logprob": -2.9050588607788086, "rank": 2, "decoded_token": " environment"}, "12097": {"logprob": -2.9675588607788086, "rank": 3, "decoded_token": " park"}, "28035": {"logprob": -3.5300588607788086, "rank": 4, "decoded_token": " landscape"}, "4457": {"logprob": -3.7175588607788086, "rank": 5, "decoded_token": " area"}}, {"1338": {"logprob": -0.40908101201057434, "rank": 1, "decoded_token": ".\n\n"}, "1626": {"logprob": -2.284080982208252, "rank": 2, "decoded_token": ".\n"}, "2425": {"logprob": -2.721580982208252, "rank": 3, "decoded_token": " under"}, "1454": {"logprob": -3.096580982208252, "rank": 4, "decoded_token": " with"}, "1044": {"logprob": -3.284080982208252, "rank": 5, "decoded_token": ","}}, {"1052": {"logprob": -1.0536651611328125, "rank": 1, "decoded_token": "4"}, "11745": {"logprob": -2.1161651611328125, "rank": 2, "decoded_token": "Here"}, "16840": {"logprob": -3.2411651611328125, "rank": 3, "decoded_token": "Since"}, "2892": {"logprob": -3.3661651611328125, "rank": 4, "decoded_token": "To"}, "1040": {"logprob": -3.3661651611328125, "rank": 5, "decoded_token": "("}}, {"1046": {"logprob": -0.15314939618110657, "rank": 1, "decoded_token": "."}, "1626": {"logprob": -4.215649604797363, "rank": 2, "decoded_token": ".\n"}, "1319": {"logprob": -4.528149604797363, "rank": 3, "decoded_token": " ("}, "26667": {"logprob": -5.840649604797363, "rank": 4, "decoded_token": ".("}, "9380": {"logprob": -6.340649604797363, "rank": 5, "decoded_token": " Here"}}, {"1531": {"logprob": -1.7034302949905396, "rank": 1, "decoded_token": " The"}, "9380": {"logprob": -2.20343017578125, "rank": 2, "decoded_token": " Here"}, "1319": {"logprob": -2.26593017578125, "rank": 3, "decoded_token": " ("}, "2898": {"logprob": -2.95343017578125, "rank": 4, "decoded_token": " For"}, "9748": {"logprob": -3.01593017578125, "rank": 5, "decoded_token": " Since"}}, {"19723": {"logprob": -1.0940117835998535, "rank": 1, "decoded_token": " fifth"}, "5662": {"logprob": -3.2190117835998535, "rank": 2, "decoded_token": " provided"}, "3804": {"logprob": -3.4690117835998535, "rank": 3, "decoded_token": " last"}, "12432": {"logprob": -3.5940117835998535, "rank": 4, "decoded_token": " fourth"}, "5719": {"logprob": -3.6565117835998535, "rank": 5, "decoded_token": " initial"}}, {"3937": {"logprob": -1.0828688144683838, "rank": 1, "decoded_token": " image"}, "1319": {"logprob": -2.020368814468384, "rank": 2, "decoded_token": " ("}, "1925": {"logprob": -2.332868814468384, "rank": 3, "decoded_token": " one"}, "2016": {"logprob": -3.457868814468384, "rank": 4, "decoded_token": " set"}, "5662": {"logprob": -3.895368814468384, "rank": 5, "decoded_token": " provided"}}, {"1319": {"logprob": -0.9146400094032288, "rank": 1, "decoded_token": " ("}, "10045": {"logprob": -2.789639949798584, "rank": 2, "decoded_token": " isn"}, "1395": {"logprob": -2.914639949798584, "rank": 3, "decoded_token": " is"}, "5662": {"logprob": -3.477139949798584, "rank": 4, "decoded_token": " provided"}, "1294": {"logprob": -3.602139949798584, "rank": 5, "decoded_token": " in"}}, {"3715": {"logprob": -2.2420246601104736, "rank": 1, "decoded_token": "not"}, "3265": {"logprob": -3.3045246601104736, "rank": 2, "decoded_token": "the"}, "11018": {"logprob": -3.3670246601104736, "rank": 3, "decoded_token": "which"}, "1391": {"logprob": -3.4920246601104736, "rank": 4, "decoded_token": "if"}, "5011": {"logprob": -3.5545246601104736, "rank": 5, "decoded_token": "from"}}, {"24512": {"logprob": -1.8148819208145142, "rank": 1, "decoded_token": " applicable"}, "4326": {"logprob": -2.6273818016052246, "rank": 2, "decoded_token": " shown"}, "5662": {"logprob": -2.8773818016052246, "rank": 3, "decoded_token": " provided"}, "5656": {"logprob": -3.1898818016052246, "rank": 4, "decoded_token": " included"}, "1805": {"logprob": -3.1898818016052246, "rank": 5, "decoded_token": " part"}}, {"1435": {"logprob": -1.9905282258987427, "rank": 1, "decoded_token": " as"}, "1317": {"logprob": -2.115528106689453, "rank": 2, "decoded_token": " to"}, "1294": {"logprob": -2.303028106689453, "rank": 3, "decoded_token": " in"}, "3226": {"logprob": -2.553028106689453, "rank": 4, "decoded_token": " here"}, "1394": {"logprob": -2.553028106689453, "rank": 5, "decoded_token": " for"}}, {"1651": {"logprob": -1.3723630905151367, "rank": 1, "decoded_token": " per"}, "1494": {"logprob": -1.6848630905151367, "rank": 2, "decoded_token": " it"}, "2156": {"logprob": -2.4348630905151367, "rank": 3, "decoded_token": " there"}, "1278": {"logprob": -2.9973630905151367, "rank": 4, "decoded_token": " the"}, "1636": {"logprob": -3.0598630905151367, "rank": 5, "decoded_token": " you"}}, {"1278": {"logprob": -1.4684513807296753, "rank": 1, "decoded_token": " the"}, "2143": {"logprob": -1.5309513807296753, "rank": 2, "decoded_token": " your"}, "5719": {"logprob": -2.218451499938965, "rank": 3, "decoded_token": " initial"}, "5662": {"logprob": -3.280951499938965, "rank": 4, "decoded_token": " provided"}, "4618": {"logprob": -3.405951499938965, "rank": 5, "decoded_token": " original"}}, {"5719": {"logprob": -1.7317447662353516, "rank": 1, "decoded_token": " initial"}, "5662": {"logprob": -2.4817447662353516, "rank": 2, "decoded_token": " provided"}, "4618": {"logprob": -2.6692447662353516, "rank": 3, "decoded_token": " original"}, "2158": {"logprob": -3.1067447662353516, "rank": 4, "decoded_token": " first"}, "4265": {"logprob": -3.2317447662353516, "rank": 5, "decoded_token": " given"}}, {"4546": {"logprob": -2.323227643966675, "rank": 1, "decoded_token": " request"}, "4098": {"logprob": -2.635727643966675, "rank": 2, "decoded_token": " question"}, "5662": {"logprob": -2.760727643966675, "rank": 3, "decoded_token": " provided"}, "6360": {"logprob": -2.885727643966675, "rank": 4, "decoded_token": " description"}, "4249": {"logprob": -3.260727643966675, "rank": 5, "decoded_token": " single"}}, {"1041": {"logprob": -2.016270160675049, "rank": 1, "decoded_token": ")"}, "1394": {"logprob": -2.203770160675049, "rank": 2, "decoded_token": " for"}, "1044": {"logprob": -2.891270160675049, "rank": 3, "decoded_token": ","}, "3640": {"logprob": -3.016270160675049, "rank": 4, "decoded_token": "):\n"}, "1321": {"logprob": -3.078770160675049, "rank": 5, "decoded_token": " and"}}, {"2168": {"logprob": -2.4032278060913086, "rank": 1, "decoded_token": " would"}, "1395": {"logprob": -2.4657278060913086, "rank": 2, "decoded_token": " is"}, "7444": {"logprob": -2.5907278060913086, "rank": 3, "decoded_token": " seems"}, "2715": {"logprob": -2.6532278060913086, "rank": 4, "decoded_token": " should"}, "10045": {"logprob": -3.2157278060913086, "rank": 5, "decoded_token": " isn"}}, {"1402": {"logprob": -1.638601303100586, "rank": 1, "decoded_token": " be"}, "1605": {"logprob": -2.138601303100586, "rank": 2, "decoded_token": " not"}, "96594": {"logprob": -2.513601303100586, "rank": 3, "decoded_token": " logically"}, "2534": {"logprob": -2.888601303100586, "rank": 4, "decoded_token": " need"}, "79961": {"logprob": -3.201101303100586, "rank": 5, "decoded_token": " ideally"}}, {"1278": {"logprob": -2.479240894317627, "rank": 1, "decoded_token": " the"}, "1261": {"logprob": -2.854240894317627, "rank": 2, "decoded_token": " a"}, "3866": {"logprob": -3.041740894317627, "rank": 3, "decoded_token": " another"}, "57773": {"logprob": -3.479240894317627, "rank": 4, "decoded_token": " irrelevant"}, "1877": {"logprob": -3.604240894317627, "rank": 5, "decoded_token": ":\n"}}, {"1925": {"logprob": -2.5215530395507812, "rank": 1, "decoded_token": " one"}, "2879": {"logprob": -3.1465530395507812, "rank": 2, "decoded_token": " same"}, "25342": {"logprob": -3.4590530395507812, "rank": 3, "decoded_token": " incorrect"}, "7244": {"logprob": -3.5840530395507812, "rank": 4, "decoded_token": " black"}, "4275": {"logprob": -3.5840530395507812, "rank": 5, "decoded_token": " next"}}, {"1454": {"logprob": -2.0894155502319336, "rank": 1, "decoded_token": " with"}, "1636": {"logprob": -3.0269155502319336, "rank": 2, "decoded_token": " you"}, "1307": {"logprob": -3.0269155502319336, "rank": 3, "decoded_token": " of"}, "1455": {"logprob": -3.2769155502319336, "rank": 4, "decoded_token": " that"}, "1562": {"logprob": -3.4019155502319336, "rank": 5, "decoded_token": " from"}}, {"1278": {"logprob": -1.3837254047393799, "rank": 1, "decoded_token": " the"}, "1261": {"logprob": -1.6962254047393799, "rank": 2, "decoded_token": " a"}, "3866": {"logprob": -3.13372540473938, "rank": 3, "decoded_token": " another"}, "1420": {"logprob": -3.19622540473938, "rank": 4, "decoded_token": " an"}, "2586": {"logprob": -3.69622540473938, "rank": 5, "decoded_token": " just"}}, {"10575": {"logprob": -1.6755551099777222, "rank": 1, "decoded_token": " dog"}, "7244": {"logprob": -2.3005552291870117, "rank": 2, "decoded_token": " black"}, "25342": {"logprob": -3.7380552291870117, "rank": 3, "decoded_token": " incorrect"}, "3549": {"logprob": -3.8630552291870117, "rank": 4, "decoded_token": " path"}, "4811": {"logprob": -4.238055229187012, "rank": 5, "decoded_token": " specific"}}, {"2790": {"logprob": -1.985573649406433, "rank": 1, "decoded_token": " again"}, "1044": {"logprob": -2.4230737686157227, "rank": 2, "decoded_token": ","}, "1562": {"logprob": -2.7355737686157227, "rank": 3, "decoded_token": " from"}, "1338": {"logprob": -2.7980737686157227, "rank": 4, "decoded_token": ".\n\n"}, "1408": {"logprob": -3.0480737686157227, "rank": 5, "decoded_token": " on"}}, {"1044": {"logprob": -1.7941354513168335, "rank": 1, "decoded_token": ","}, "1693": {"logprob": -2.231635570526123, "rank": 2, "decoded_token": " if"}, "1809": {"logprob": -2.731635570526123, "rank": 3, "decoded_token": " but"}, "1505": {"logprob": -3.356635570526123, "rank": 4, "decoded_token": " or"}, "1562": {"logprob": -3.419135570526123, "rank": 5, "decoded_token": " from"}}, {"1809": {"logprob": -1.9572563171386719, "rank": 1, "decoded_token": " but"}, "1878": {"logprob": -2.269756317138672, "rank": 2, "decoded_token": " so"}, "1799": {"logprob": -2.769756317138672, "rank": 3, "decoded_token": " which"}, "1321": {"logprob": -3.769756317138672, "rank": 4, "decoded_token": " and"}, "1693": {"logprob": -3.957256317138672, "rank": 5, "decoded_token": " if"}}, {"4136": {"logprob": -2.412797212600708, "rank": 1, "decoded_token": " since"}, "3226": {"logprob": -2.475297212600708, "rank": 2, "decoded_token": " here"}, "1278": {"logprob": -2.850297212600708, "rank": 3, "decoded_token": " the"}, "2878": {"logprob": -2.975297212600708, "rank": 4, "decoded_token": " let"}, "1494": {"logprob": -3.287797212600708, "rank": 5, "decoded_token": " it"}}, {"1494": {"logprob": -1.4206243753433228, "rank": 1, "decoded_token": " it"}, "1278": {"logprob": -2.045624256134033, "rank": 2, "decoded_token": " the"}, "1729": {"logprob": -2.295624256134033, "rank": 3, "decoded_token": " we"}, "2156": {"logprob": -2.483124256134033, "rank": 4, "decoded_token": " there"}, "1636": {"logprob": -2.483124256134033, "rank": 5, "decoded_token": " you"}}, {"1681": {"logprob": -1.4698597192764282, "rank": 1, "decoded_token": "'s"}, "2190": {"logprob": -1.6573597192764282, "rank": 2, "decoded_token": "\u2019s"}, "1395": {"logprob": -2.1573596000671387, "rank": 3, "decoded_token": " is"}, "1486": {"logprob": -2.5323596000671387, "rank": 4, "decoded_token": " was"}, "7444": {"logprob": -3.2823596000671387, "rank": 5, "decoded_token": " seems"}}, {"5314": {"logprob": -1.6933988332748413, "rank": 1, "decoded_token": " already"}, "1605": {"logprob": -2.130898952484131, "rank": 2, "decoded_token": " not"}, "1278": {"logprob": -2.193398952484131, "rank": 3, "decoded_token": " the"}, "2342": {"logprob": -3.318398952484131, "rank": 4, "decoded_token": " only"}, "13578": {"logprob": -3.568398952484131, "rank": 5, "decoded_token": " repeated"}}, {"5055": {"logprob": -1.3677805662155151, "rank": 1, "decoded_token": " described"}, "13875": {"logprob": -1.4927805662155151, "rank": 2, "decoded_token": " covered"}, "24511": {"logprob": -2.6802806854248047, "rank": 3, "decoded_token": " addressed"}, "10910": {"logprob": -2.7427806854248047, "rank": 4, "decoded_token": " mentioned"}, "5656": {"logprob": -3.4302806854248047, "rank": 5, "decoded_token": " included"}}, {"1044": {"logprob": -1.6195299625396729, "rank": 1, "decoded_token": ","}, "2100": {"logprob": -2.057029962539673, "rank": 2, "decoded_token": ":\n\n"}, "1294": {"logprob": -2.057029962539673, "rank": 3, "decoded_token": " in"}, "1877": {"logprob": -2.369529962539673, "rank": 4, "decoded_token": ":\n"}, "15423": {"logprob": -2.869529962539673, "rank": 5, "decoded_token": " initially"}}, {"3226": {"logprob": -1.7415276765823364, "rank": 1, "decoded_token": " here"}, "2878": {"logprob": -1.8665276765823364, "rank": 2, "decoded_token": " let"}, "1362": {"logprob": -2.304027557373047, "rank": 3, "decoded_token": " I"}, "1494": {"logprob": -2.616527557373047, "rank": 4, "decoded_token": " it"}, "1278": {"logprob": -3.179027557373047, "rank": 5, "decoded_token": " the"}}, {"2190": {"logprob": -1.3488249778747559, "rank": 1, "decoded_token": "\u2019s"}, "1681": {"logprob": -1.4113249778747559, "rank": 2, "decoded_token": "'s"}, "1395": {"logprob": -1.5363249778747559, "rank": 3, "decoded_token": " is"}, "1584": {"logprob": -2.473824977874756, "rank": 4, "decoded_token": " are"}, "1494": {"logprob": -4.348824977874756, "rank": 5, "decoded_token": " it"}}, {"1261": {"logprob": -1.4024583101272583, "rank": 1, "decoded_token": " a"}, "1278": {"logprob": -1.9024583101272583, "rank": 2, "decoded_token": " the"}, "3866": {"logprob": -2.1524581909179688, "rank": 3, "decoded_token": " another"}, "1420": {"logprob": -2.5274581909179688, "rank": 4, "decoded_token": " an"}, "1605": {"logprob": -4.152458190917969, "rank": 5, "decoded_token": " not"}}, {"2801": {"logprob": -2.7750463485717773, "rank": 1, "decoded_token": " different"}, "6468": {"logprob": -3.1500463485717773, "rank": 2, "decoded_token": " focus"}, "21788": {"logprob": -3.2750463485717773, "rank": 3, "decoded_token": " repeat"}, "17793": {"logprob": -3.4625463485717773, "rank": 4, "decoded_token": " summary"}, "13426": {"logprob": -3.6500463485717773, "rank": 5, "decoded_token": " brief"}}, {"6468": {"logprob": -2.8590242862701416, "rank": 1, "decoded_token": " focus"}, "1925": {"logprob": -3.4215242862701416, "rank": 2, "decoded_token": " one"}, "19190": {"logprob": -3.9215242862701416, "rank": 3, "decoded_token": " interpretation"}, "6360": {"logprob": -3.9840242862701416, "rank": 4, "decoded_token": " description"}, "3336": {"logprob": -3.9840242862701416, "rank": 5, "decoded_token": " example"}}, {"1693": {"logprob": -2.391437530517578, "rank": 1, "decoded_token": " if"}, "1877": {"logprob": -2.516437530517578, "rank": 2, "decoded_token": ":\n"}, "1408": {"logprob": -2.516437530517578, "rank": 3, "decoded_token": " on"}, "1058": {"logprob": -2.578937530517578, "rank": 4, "decoded_token": ":"}, "2100": {"logprob": -2.703937530517578, "rank": 5, "decoded_token": ":\n\n"}}, {"1729": {"logprob": -2.6488425731658936, "rank": 1, "decoded_token": " we"}, "6618": {"logprob": -2.6488425731658936, "rank": 2, "decoded_token": " needed"}, "1494": {"logprob": -2.9613425731658936, "rank": 3, "decoded_token": " it"}, "2258": {"logprob": -3.2113425731658936, "rank": 4, "decoded_token": " any"}, "1636": {"logprob": -3.2738425731658936, "rank": 5, "decoded_token": " you"}}, {"3369": {"logprob": -2.0612776279449463, "rank": 1, "decoded_token": " consider"}, "1722": {"logprob": -2.6862776279449463, "rank": 2, "decoded_token": " were"}, "1880": {"logprob": -3.1237776279449463, "rank": 3, "decoded_token": " had"}, "14649": {"logprob": -3.1862776279449463, "rank": 4, "decoded_token": " assume"}, "55328": {"logprob": -3.2487776279449463, "rank": 5, "decoded_token": " mistaken"}}, {"1278": {"logprob": -1.7770118713378906, "rank": 1, "decoded_token": " the"}, "3866": {"logprob": -2.4645118713378906, "rank": 2, "decoded_token": " another"}, "1261": {"logprob": -2.5895118713378906, "rank": 3, "decoded_token": " a"}, "1494": {"logprob": -3.5270118713378906, "rank": 4, "decoded_token": " it"}, "1420": {"logprob": -3.5895118713378906, "rank": 5, "decoded_token": " an"}}, {"3629": {"logprob": -3.3516223430633545, "rank": 1, "decoded_token": " following"}, "4275": {"logprob": -3.5391223430633545, "rank": 2, "decoded_token": " next"}, "12432": {"logprob": -3.7266223430633545, "rank": 3, "decoded_token": " fourth"}, "5719": {"logprob": -3.7891223430633545, "rank": 4, "decoded_token": " initial"}, "2158": {"logprob": -3.9141223430633545, "rank": 5, "decoded_token": " first"}}, {"75275": {"logprob": -3.0541036128997803, "rank": 1, "decoded_token": " hypothetical"}, "2100": {"logprob": -3.1166036128997803, "rank": 2, "decoded_token": ":\n\n"}, "1877": {"logprob": -3.1166036128997803, "rank": 3, "decoded_token": ":\n"}, "1435": {"logprob": -4.179103851318359, "rank": 4, "decoded_token": " as"}, "1319": {"logprob": -4.179103851318359, "rank": 5, "decoded_token": " ("}}, {"1877": {"logprob": -2.70977520942688, "rank": 1, "decoded_token": ":\n"}, "2100": {"logprob": -2.83477520942688, "rank": 2, "decoded_token": ":\n\n"}, "1115": {"logprob": -3.20977520942688, "rank": 3, "decoded_token": "s"}, "12432": {"logprob": -3.64727520942688, "rank": 4, "decoded_token": " fourth"}, "1058": {"logprob": -3.89727520942688, "rank": 5, "decoded_token": ":"}}, {"1256": {"logprob": -0.4959573745727539, "rank": 1, "decoded_token": " "}, "1293": {"logprob": -1.870957374572754, "rank": 2, "decoded_token": " "}, "11745": {"logprob": -4.183457374572754, "rank": 3, "decoded_token": "Here"}, "1784": {"logprob": -4.495957374572754, "rank": 4, "decoded_token": "The"}, "1053": {"logprob": -4.495957374572754, "rank": 5, "decoded_token": "5"}}, {"1462": {"logprob": -1.5276975631713867, "rank": 1, "decoded_token": " -"}, "1319": {"logprob": -2.4026975631713867, "rank": 2, "decoded_token": " ("}, "9246": {"logprob": -3.2776975631713867, "rank": 3, "decoded_token": " Let"}, "9380": {"logprob": -3.3401975631713867, "rank": 4, "decoded_token": " Here"}, "1531": {"logprob": -3.3401975631713867, "rank": 5, "decoded_token": " The"}}, {"1531": {"logprob": -2.32918381690979, "rank": 1, "decoded_token": " The"}, "1319": {"logprob": -2.82918381690979, "rank": 2, "decoded_token": " ("}, "1349": {"logprob": -3.01668381690979, "rank": 3, "decoded_token": " A"}, "2898": {"logprob": -3.20418381690979, "rank": 4, "decoded_token": " For"}, "9380": {"logprob": -3.57918381690979, "rank": 5, "decoded_token": " Here"}}, {"5719": {"logprob": -2.6479814052581787, "rank": 1, "decoded_token": " initial"}, "19723": {"logprob": -3.2729814052581787, "rank": 2, "decoded_token": " fifth"}, "3804": {"logprob": -3.3979814052581787, "rank": 3, "decoded_token": " last"}, "25342": {"logprob": -3.8979814052581787, "rank": 4, "decoded_token": " incorrect"}, "5662": {"logprob": -3.9604814052581787, "rank": 5, "decoded_token": " provided"}}, {"5662": {"logprob": -2.9346542358398438, "rank": 1, "decoded_token": " provided"}, "6468": {"logprob": -3.1846542358398438, "rank": 2, "decoded_token": " focus"}, "7244": {"logprob": -3.6221542358398438, "rank": 3, "decoded_token": " black"}, "10575": {"logprob": -3.6221542358398438, "rank": 4, "decoded_token": " dog"}, "6360": {"logprob": -3.8096542358398438, "rank": 5, "decoded_token": " description"}}, {"3937": {"logprob": -3.0049667358398438, "rank": 1, "decoded_token": " image"}, "10575": {"logprob": -3.0674667358398438, "rank": 2, "decoded_token": " dog"}, "8061": {"logprob": -3.5049667358398438, "rank": 3, "decoded_token": " images"}, "2016": {"logprob": -3.5049667358398438, "rank": 4, "decoded_token": " set"}, "2667": {"logprob": -3.5674667358398438, "rank": 5, "decoded_token": " second"}}, {"1395": {"logprob": -2.7246885299682617, "rank": 1, "decoded_token": " is"}, "1319": {"logprob": -2.7871885299682617, "rank": 2, "decoded_token": " ("}, "1307": {"logprob": -3.6621885299682617, "rank": 3, "decoded_token": " of"}, "2016": {"logprob": -3.7871885299682617, "rank": 4, "decoded_token": " set"}, "1058": {"logprob": -3.7871885299682617, "rank": 5, "decoded_token": ":"}}, {"1261": {"logprob": -2.914872646331787, "rank": 1, "decoded_token": " a"}, "1605": {"logprob": -3.164872646331787, "rank": 2, "decoded_token": " not"}, "1278": {"logprob": -3.164872646331787, "rank": 3, "decoded_token": " the"}, "2586": {"logprob": -3.227372646331787, "rank": 4, "decoded_token": " just"}, "13578": {"logprob": -3.352372646331787, "rank": 5, "decoded_token": " repeated"}}, {"7244": {"logprob": -2.4122366905212402, "rank": 1, "decoded_token": " black"}, "15115": {"logprob": -3.1622366905212402, "rank": 2, "decoded_token": " detailed"}, "10575": {"logprob": -3.1622366905212402, "rank": 3, "decoded_token": " dog"}, "6468": {"logprob": -3.2247366905212402, "rank": 4, "decoded_token": " focus"}, "6165": {"logprob": -3.3497366905212402, "rank": 5, "decoded_token": " simple"}}, {"10575": {"logprob": -0.2919383645057678, "rank": 1, "decoded_token": " dog"}, "119075": {"logprob": -3.104438304901123, "rank": 2, "decoded_token": " Labrador"}, "116572": {"logprob": -3.479438304901123, "rank": 3, "decoded_token": " puppy"}, "15812": {"logprob": -4.104438304901123, "rank": 4, "decoded_token": " Lab"}, "2001": {"logprob": -4.479438304901123, "rank": 5, "decoded_token": " po"}}, {"7283": {"logprob": -1.594571590423584, "rank": 1, "decoded_token": " looking"}, "1408": {"logprob": -2.219571590423584, "rank": 2, "decoded_token": " on"}, "18970": {"logprob": -2.407071590423584, "rank": 3, "decoded_token": " sitting"}, "1454": {"logprob": -2.657071590423584, "rank": 4, "decoded_token": " with"}, "11589": {"logprob": -3.282071590423584, "rank": 5, "decoded_token": " gaz"}}, {"2015": {"logprob": -2.1110990047454834, "rank": 1, "decoded_token": " up"}, "7655": {"logprob": -2.1735990047454834, "rank": 2, "decoded_token": " directly"}, "1935": {"logprob": -2.4860990047454834, "rank": 3, "decoded_token": " int"}, "4524": {"logprob": -3.0485990047454834, "rank": 4, "decoded_token": " cur"}, "40022": {"logprob": -3.2985990047454834, "rank": 5, "decoded_token": " upward"}}, {"1454": {"logprob": -1.676759123802185, "rank": 1, "decoded_token": " with"}, "1513": {"logprob": -2.1142592430114746, "rank": 2, "decoded_token": " at"}, "4914": {"logprob": -2.5517592430114746, "rank": 3, "decoded_token": " thought"}, "1935": {"logprob": -2.8642592430114746, "rank": 4, "decoded_token": " int"}, "1408": {"logprob": -2.8642592430114746, "rank": 5, "decoded_token": " on"}}, {"1261": {"logprob": -1.243793249130249, "rank": 1, "decoded_token": " a"}, "1420": {"logprob": -1.931293249130249, "rank": 2, "decoded_token": " an"}, "2246": {"logprob": -2.618793249130249, "rank": 3, "decoded_token": " its"}, "9924": {"logprob": -2.618793249130249, "rank": 4, "decoded_token": " wide"}, "14781": {"logprob": -3.743793249130249, "rank": 5, "decoded_token": " focused"}}, {"26517": {"logprob": -1.6467483043670654, "rank": 1, "decoded_token": " calm"}, "14781": {"logprob": -2.5842483043670654, "rank": 2, "decoded_token": " focused"}, "29691": {"logprob": -2.7092483043670654, "rank": 3, "decoded_token": " contempl"}, "85596": {"logprob": -3.1467483043670654, "rank": 4, "decoded_token": " solemn"}, "16318": {"logprob": -3.2092483043670654, "rank": 5, "decoded_token": " neutral"}}, {"4818": {"logprob": -0.6443419456481934, "rank": 1, "decoded_token": " expression"}, "1311": {"logprob": -2.2068419456481934, "rank": 2, "decoded_token": " de"}, "1321": {"logprob": -2.5193419456481934, "rank": 3, "decoded_token": " and"}, "1044": {"logprob": -2.9568419456481934, "rank": 4, "decoded_token": ","}, "22131": {"logprob": -3.4568419456481934, "rank": 5, "decoded_token": " gaze"}}, {"1408": {"logprob": -0.7509063482284546, "rank": 1, "decoded_token": " on"}, "1338": {"logprob": -2.188406467437744, "rank": 2, "decoded_token": ".\n\n"}, "2136": {"logprob": -2.875906467437744, "rank": 3, "decoded_token": " over"}, "3675": {"logprob": -3.188406467437744, "rank": 4, "decoded_token": " against"}, "1046": {"logprob": -3.313406467437744, "rank": 5, "decoded_token": "."}}, {"1261": {"logprob": -0.41352635622024536, "rank": 1, "decoded_token": " a"}, "44130": {"logprob": -2.2885262966156006, "rank": 2, "decoded_token": " rust"}, "32656": {"logprob": -2.8510262966156006, "rank": 3, "decoded_token": " wooden"}, "2549": {"logprob": -3.6010262966156006, "rank": 4, "decoded_token": " what"}, "3403": {"logprob": -4.03852653503418, "rank": 5, "decoded_token": " text"}}, {"44130": {"logprob": -1.1644353866577148, "rank": 1, "decoded_token": " rust"}, "3403": {"logprob": -1.7269353866577148, "rank": 2, "decoded_token": " text"}, "32656": {"logprob": -1.7894353866577148, "rank": 3, "decoded_token": " wooden"}, "6165": {"logprob": -3.351935386657715, "rank": 4, "decoded_token": " simple"}, "1615": {"logprob": -3.601935386657715, "rank": 5, "decoded_token": " pl"}}, {"1290": {"logprob": -0.007089222315698862, "rank": 1, "decoded_token": "ic"}, "1121": {"logprob": -6.319589138031006, "rank": 2, "decoded_token": "y"}, "2981": {"logprob": -6.507089138031006, "rank": 3, "decoded_token": "ically"}, "12500": {"logprob": -6.694589138031006, "rank": 4, "decoded_token": "icated"}, "86794": {"logprob": -7.632089138031006, "rank": 5, "decoded_token": "-colored"}}, {"32656": {"logprob": -0.5749364495277405, "rank": 1, "decoded_token": " wooden"}, "1615": {"logprob": -2.4499363899230957, "rank": 2, "decoded_token": " pl"}, "1044": {"logprob": -2.6999363899230957, "rank": 3, "decoded_token": ","}, "12603": {"logprob": -2.8249363899230957, "rank": 4, "decoded_token": " wood"}, "4691": {"logprob": -3.3249363899230957, "rank": 5, "decoded_token": " surface"}}, {"1615": {"logprob": -0.6836710572242737, "rank": 1, "decoded_token": " pl"}, "7042": {"logprob": -2.433670997619629, "rank": 2, "decoded_token": " background"}, "4691": {"logprob": -2.433670997619629, "rank": 3, "decoded_token": " surface"}, "92504": {"logprob": -2.558670997619629, "rank": 4, "decoded_token": " backdrop"}, "9710": {"logprob": -3.058670997619629, "rank": 5, "decoded_token": " board"}}, {"2395": {"logprob": -0.063260018825531, "rank": 1, "decoded_token": "ank"}, "5933": {"logprob": -3.063260078430176, "rank": 2, "decoded_token": "anks"}, "122370": {"logprob": -4.313260078430176, "rank": 3, "decoded_token": "anking"}, "4713": {"logprob": -7.688260078430176, "rank": 4, "decoded_token": "atter"}, "58739": {"logprob": -8.563260078430176, "rank": 5, "decoded_token": "ANK"}}, {"7042": {"logprob": -1.4006984233856201, "rank": 1, "decoded_token": " background"}, "92504": {"logprob": -1.6506984233856201, "rank": 2, "decoded_token": " backdrop"}, "4691": {"logprob": -2.08819842338562, "rank": 3, "decoded_token": " surface"}, "1338": {"logprob": -2.83819842338562, "rank": 4, "decoded_token": ".\n\n"}, "9436": {"logprob": -2.96319842338562, "rank": 5, "decoded_token": " setting"}}, {"1338": {"logprob": -0.7081566452980042, "rank": 1, "decoded_token": ".\n\n"}, "1046": {"logprob": -0.8956566452980042, "rank": 2, "decoded_token": "."}, "1319": {"logprob": -4.145656585693359, "rank": 3, "decoded_token": " ("}, "2100": {"logprob": -4.395656585693359, "rank": 4, "decoded_token": ":\n\n"}, "1626": {"logprob": -4.708156585693359, "rank": 5, "decoded_token": ".\n"}}, {"2892": {"logprob": -2.5741822719573975, "rank": 1, "decoded_token": "To"}, "4393": {"logprob": -2.8866822719573975, "rank": 2, "decoded_token": "For"}, "11745": {"logprob": -3.1366822719573975, "rank": 3, "decoded_token": "Here"}, "12598": {"logprob": -3.1991822719573975, "rank": 4, "decoded_token": "Let"}, "1040": {"logprob": -3.2616822719573975, "rank": 5, "decoded_token": "("}}, {"38695": {"logprob": -1.3182454109191895, "rank": 1, "decoded_token": " clarify"}, "10035": {"logprob": -2.3807454109191895, "rank": 2, "decoded_token": " avoid"}, "11811": {"logprob": -3.0057454109191895, "rank": 3, "decoded_token": " ensure"}, "66370": {"logprob": -3.0682454109191895, "rank": 4, "decoded_token": " summarize"}, "36993": {"logprob": -3.4432454109191895, "rank": 5, "decoded_token": " strictly"}}, {"1044": {"logprob": -1.8413705825805664, "rank": 1, "decoded_token": ","}, "1278": {"logprob": -2.5913705825805664, "rank": 2, "decoded_token": " the"}, "1321": {"logprob": -2.9038705825805664, "rank": 3, "decoded_token": " and"}, "1877": {"logprob": -3.2163705825805664, "rank": 4, "decoded_token": ":\n"}, "4057": {"logprob": -3.5288705825805664, "rank": 5, "decoded_token": " based"}}, {"3226": {"logprob": -1.5737794637680054, "rank": 1, "decoded_token": " here"}, "1362": {"logprob": -2.698779582977295, "rank": 2, "decoded_token": " I"}, "1278": {"logprob": -2.761279582977295, "rank": 3, "decoded_token": " the"}, "2342": {"logprob": -2.886279582977295, "rank": 4, "decoded_token": " only"}, "2878": {"logprob": -2.948779582977295, "rank": 5, "decoded_token": " let"}}, {"2190": {"logprob": -1.2323426008224487, "rank": 1, "decoded_token": "\u2019s"}, "1681": {"logprob": -1.5448426008224487, "rank": 2, "decoded_token": "'s"}, "1584": {"logprob": -1.5448426008224487, "rank": 3, "decoded_token": " are"}, "1395": {"logprob": -2.0448427200317383, "rank": 4, "decoded_token": " is"}, "6510": {"logprob": -4.607342720031738, "rank": 5, "decoded_token": "with"}}, {"1278": {"logprob": -1.6423555612564087, "rank": 1, "decoded_token": " the"}, "1261": {"logprob": -1.7048555612564087, "rank": 2, "decoded_token": " a"}, "2606": {"logprob": -3.267355442047119, "rank": 3, "decoded_token": " how"}, "1420": {"logprob": -3.329855442047119, "rank": 4, "decoded_token": " an"}, "2342": {"logprob": -3.517355442047119, "rank": 5, "decoded_token": " only"}}, {"6298": {"logprob": -2.6795427799224854, "rank": 1, "decoded_token": " correct"}, "4811": {"logprob": -3.3670427799224854, "rank": 2, "decoded_token": " specific"}, "17793": {"logprob": -3.6795427799224854, "rank": 3, "decoded_token": " summary"}, "15115": {"logprob": -3.8045427799224854, "rank": 4, "decoded_token": " detailed"}, "6468": {"logprob": -3.8670427799224854, "rank": 5, "decoded_token": " focus"}}, {"6360": {"logprob": -2.9291810989379883, "rank": 1, "decoded_token": " description"}, "6468": {"logprob": -3.4291810989379883, "rank": 2, "decoded_token": " focus"}, "7980": {"logprob": -3.5541810989379883, "rank": 3, "decoded_token": " sequence"}, "1321": {"logprob": -3.5541810989379883, "rank": 4, "decoded_token": " and"}, "44433": {"logprob": -3.8041810989379883, "rank": 5, "decoded_token": " breakdown"}}, {"1394": {"logprob": -1.4118860960006714, "rank": 1, "decoded_token": " for"}, "1307": {"logprob": -2.161886215209961, "rank": 2, "decoded_token": " of"}, "4057": {"logprob": -2.599386215209961, "rank": 3, "decoded_token": " based"}, "30557": {"logprob": -3.724386215209961, "rank": 4, "decoded_token": " focusing"}, "2342": {"logprob": -3.911886215209961, "rank": 5, "decoded_token": " only"}}, {"1278": {"logprob": -1.170920729637146, "rank": 1, "decoded_token": " the"}, "2744": {"logprob": -2.4834208488464355, "rank": 2, "decoded_token": " each"}, "2143": {"logprob": -3.1709208488464355, "rank": 3, "decoded_token": " your"}, "1747": {"logprob": -3.2959208488464355, "rank": 4, "decoded_token": " all"}, "2342": {"logprob": -3.3584208488464355, "rank": 5, "decoded_token": " only"}}, {"5662": {"logprob": -2.8693411350250244, "rank": 1, "decoded_token": " provided"}, "5719": {"logprob": -3.1193411350250244, "rank": 2, "decoded_token": " initial"}, "8061": {"logprob": -3.3693411350250244, "rank": 3, "decoded_token": " images"}, "12432": {"logprob": -3.4318411350250244, "rank": 4, "decoded_token": " fourth"}, "2667": {"logprob": -3.4943411350250244, "rank": 5, "decoded_token": " second"}}, {"7244": {"logprob": -2.477940559387207, "rank": 1, "decoded_token": " black"}, "10575": {"logprob": -2.727940559387207, "rank": 2, "decoded_token": " dog"}, "3937": {"logprob": -2.727940559387207, "rank": 3, "decoded_token": " image"}, "8061": {"logprob": -2.790440559387207, "rank": 4, "decoded_token": " images"}, "2016": {"logprob": -2.852940559387207, "rank": 5, "decoded_token": " set"}}, {"10575": {"logprob": -0.09506329894065857, "rank": 1, "decoded_token": " dog"}, "3937": {"logprob": -3.8450632095336914, "rank": 2, "decoded_token": " image"}, "63524": {"logprob": -4.470063209533691, "rank": 3, "decoded_token": "dog"}, "3028": {"logprob": -4.970063209533691, "rank": 4, "decoded_token": "-d"}, "1321": {"logprob": -5.220063209533691, "rank": 5, "decoded_token": " and"}}, {"3937": {"logprob": -0.6087309122085571, "rank": 1, "decoded_token": " image"}, "2100": {"logprob": -2.7337307929992676, "rank": 2, "decoded_token": ":\n\n"}, "1294": {"logprob": -3.2337307929992676, "rank": 3, "decoded_token": " in"}, "13083": {"logprob": -3.3587307929992676, "rank": 4, "decoded_token": " picture"}, "1877": {"logprob": -3.3587307929992676, "rank": 5, "decoded_token": ":\n"}}, {"1877": {"logprob": -1.969537615776062, "rank": 1, "decoded_token": ":\n"}, "2100": {"logprob": -2.0945377349853516, "rank": 2, "decoded_token": ":\n\n"}, "2342": {"logprob": -2.4695377349853516, "rank": 3, "decoded_token": " only"}, "13703": {"logprob": -3.0945377349853516, "rank": 4, "decoded_token": " specifically"}, "9412": {"logprob": -3.1570377349853516, "rank": 5, "decoded_token": " alone"}}, {"1065": {"logprob": -1.341880202293396, "rank": 1, "decoded_token": "A"}, "1784": {"logprob": -1.841880202293396, "rank": 2, "decoded_token": "The"}, "1256": {"logprob": -1.966880202293396, "rank": 3, "decoded_token": " "}, "1438": {"logprob": -2.4668803215026855, "rank": 4, "decoded_token": "**"}, "1045": {"logprob": -2.5293803215026855, "rank": 5, "decoded_token": "-"}}, {"7244": {"logprob": -0.6517431139945984, "rank": 1, "decoded_token": " black"}, "6231": {"logprob": -2.714243173599243, "rank": 2, "decoded_token": " close"}, "85596": {"logprob": -3.026743173599243, "rank": 3, "decoded_token": " solemn"}, "14781": {"logprob": -3.401743173599243, "rank": 4, "decoded_token": " focused"}, "26517": {"logprob": -3.526743173599243, "rank": 5, "decoded_token": " calm"}}, {"10575": {"logprob": -0.10892026871442795, "rank": 1, "decoded_token": " dog"}, "119075": {"logprob": -3.7964203357696533, "rank": 2, "decoded_token": " Labrador"}, "116572": {"logprob": -3.9839203357696533, "rank": 3, "decoded_token": " puppy"}, "1044": {"logprob": -4.983920097351074, "rank": 4, "decoded_token": ","}, "94057": {"logprob": -5.171420097351074, "rank": 5, "decoded_token": " canine"}}, {"11589": {"logprob": -1.534816861152649, "rank": 1, "decoded_token": " gaz"}, "1395": {"logprob": -1.722316861152649, "rank": 2, "decoded_token": " is"}, "1454": {"logprob": -2.0348167419433594, "rank": 3, "decoded_token": " with"}, "53048": {"logprob": -2.2848167419433594, "rank": 4, "decoded_token": " sits"}, "10637": {"logprob": -2.6598167419433594, "rank": 5, "decoded_token": " looks"}}, {"1264": {"logprob": -0.4754399061203003, "rank": 1, "decoded_token": "es"}, "1302": {"logprob": -0.9754399061203003, "rank": 2, "decoded_token": "ing"}, "1944": {"logprob": -8.16294002532959, "rank": 3, "decoded_token": "ely"}, "47885": {"logprob": -8.16294002532959, "rank": 4, "decoded_token": "edly"}, "15006": {"logprob": -8.41294002532959, "rank": 5, "decoded_token": "ingly"}}, {"1935": {"logprob": -1.7025537490844727, "rank": 1, "decoded_token": " int"}, "7655": {"logprob": -2.0775537490844727, "rank": 2, "decoded_token": " directly"}, "40022": {"logprob": -2.2650537490844727, "rank": 3, "decoded_token": " upward"}, "2015": {"logprob": -2.4525537490844727, "rank": 4, "decoded_token": " up"}, "74606": {"logprob": -2.7025537490844727, "rank": 5, "decoded_token": " upwards"}}, {"3929": {"logprob": -0.0234219990670681, "rank": 1, "decoded_token": "ently"}, "2749": {"logprob": -5.335921764373779, "rank": 2, "decoded_token": "ros"}, "1533": {"logprob": -6.398421764373779, "rank": 3, "decoded_token": "ang"}, "3923": {"logprob": -6.523421764373779, "rank": 4, "decoded_token": "ively"}, "20626": {"logprob": -6.773421764373779, "rank": 5, "decoded_token": "rep"}}, {"40022": {"logprob": -0.8856239318847656, "rank": 1, "decoded_token": " upward"}, "74606": {"logprob": -1.5106239318847656, "rank": 2, "decoded_token": " upwards"}, "1454": {"logprob": -2.2606239318847656, "rank": 3, "decoded_token": " with"}, "8848": {"logprob": -2.7606239318847656, "rank": 4, "decoded_token": " forward"}, "1513": {"logprob": -3.5106239318847656, "rank": 5, "decoded_token": " at"}}, {"1454": {"logprob": -1.170392632484436, "rank": 1, "decoded_token": " with"}, "1408": {"logprob": -1.295392632484436, "rank": 2, "decoded_token": " on"}, "1562": {"logprob": -2.1703925132751465, "rank": 3, "decoded_token": " from"}, "3016": {"logprob": -2.4828925132751465, "rank": 4, "decoded_token": " while"}, "3675": {"logprob": -3.1703925132751465, "rank": 5, "decoded_token": " against"}}, {"1261": {"logprob": -0.8887192010879517, "rank": 1, "decoded_token": " a"}, "1420": {"logprob": -1.7637192010879517, "rank": 2, "decoded_token": " an"}, "2246": {"logprob": -2.138719081878662, "rank": 3, "decoded_token": " its"}, "9924": {"logprob": -2.888719081878662, "rank": 4, "decoded_token": " wide"}, "26517": {"logprob": -3.888719081878662, "rank": 5, "decoded_token": " calm"}}, {"26517": {"logprob": -1.2183846235275269, "rank": 1, "decoded_token": " calm"}, "2169": {"logprob": -2.8433847427368164, "rank": 2, "decoded_token": " ser"}, "16318": {"logprob": -3.0933847427368164, "rank": 3, "decoded_token": " neutral"}, "14781": {"logprob": -3.2183847427368164, "rank": 4, "decoded_token": " focused"}, "6444": {"logprob": -3.2183847427368164, "rank": 5, "decoded_token": " soft"}}, {"4818": {"logprob": -0.4964141249656677, "rank": 1, "decoded_token": " expression"}, "1321": {"logprob": -1.9339141845703125, "rank": 2, "decoded_token": " and"}, "1311": {"logprob": -2.4339141845703125, "rank": 3, "decoded_token": " de"}, "1044": {"logprob": -3.3089141845703125, "rank": 4, "decoded_token": ","}, "2985": {"logprob": -3.4964141845703125, "rank": 5, "decoded_token": " look"}}, {"1408": {"logprob": -0.8785426616668701, "rank": 1, "decoded_token": " on"}, "1562": {"logprob": -2.06604266166687, "rank": 2, "decoded_token": " from"}, "18970": {"logprob": -2.31604266166687, "rank": 3, "decoded_token": " sitting"}, "3675": {"logprob": -2.62854266166687, "rank": 4, "decoded_token": " against"}, "38235": {"logprob": -3.06604266166687, "rank": 5, "decoded_token": " resting"}}, {"1261": {"logprob": -0.3605937659740448, "rank": 1, "decoded_token": " a"}, "32656": {"logprob": -2.923093795776367, "rank": 2, "decoded_token": " wooden"}, "17253": {"logprob": -3.110593795776367, "rank": 3, "decoded_token": " weather"}, "3403": {"logprob": -3.173093795776367, "rank": 4, "decoded_token": " text"}, "44130": {"logprob": -3.298093795776367, "rank": 5, "decoded_token": " rust"}}, {"44130": {"logprob": -1.4187138080596924, "rank": 1, "decoded_token": " rust"}, "32656": {"logprob": -1.4812138080596924, "rank": 2, "decoded_token": " wooden"}, "3403": {"logprob": -1.6687138080596924, "rank": 3, "decoded_token": " text"}, "17253": {"logprob": -2.4812138080596924, "rank": 4, "decoded_token": " weather"}, "8500": {"logprob": -4.168713569641113, "rank": 5, "decoded_token": " dark"}}, {"1290": {"logprob": -0.28350138664245605, "rank": 1, "decoded_token": "ic"}, "1970": {"logprob": -2.283501386642456, "rank": 2, "decoded_token": "led"}, "1121": {"logprob": -2.971001386642456, "rank": 3, "decoded_token": "y"}, "2981": {"logprob": -4.283501625061035, "rank": 4, "decoded_token": "ically"}, "11395": {"logprob": -4.471001625061035, "rank": 5, "decoded_token": "iced"}}, {"32656": {"logprob": -0.7727869749069214, "rank": 1, "decoded_token": " wooden"}, "1044": {"logprob": -1.4602869749069214, "rank": 2, "decoded_token": ","}, "3403": {"logprob": -2.585287094116211, "rank": 3, "decoded_token": " text"}, "1615": {"logprob": -2.835287094116211, "rank": 4, "decoded_token": " pl"}, "12603": {"logprob": -3.335287094116211, "rank": 5, "decoded_token": " wood"}}, {"1615": {"logprob": -0.7352191805839539, "rank": 1, "decoded_token": " pl"}, "4691": {"logprob": -2.1102192401885986, "rank": 2, "decoded_token": " surface"}, "9710": {"logprob": -2.6102192401885986, "rank": 3, "decoded_token": " board"}, "7042": {"logprob": -2.9852192401885986, "rank": 4, "decoded_token": " background"}, "92504": {"logprob": -3.1102192401885986, "rank": 5, "decoded_token": " backdrop"}}, {"2395": {"logprob": -0.13373544812202454, "rank": 1, "decoded_token": "ank"}, "5933": {"logprob": -2.133735418319702, "rank": 2, "decoded_token": "anks"}, "122370": {"logprob": -6.883735656738281, "rank": 3, "decoded_token": "anking"}, "11847": {"logprob": -8.071235656738281, "rank": 4, "decoded_token": "anned"}, "2077": {"logprob": -8.071235656738281, "rank": 5, "decoded_token": "ink"}}, {"7042": {"logprob": -1.197163701057434, "rank": 1, "decoded_token": " background"}, "4691": {"logprob": -1.572163701057434, "rank": 2, "decoded_token": " surface"}, "92504": {"logprob": -1.884663701057434, "rank": 3, "decoded_token": " backdrop"}, "26228": {"logprob": -3.3221635818481445, "rank": 4, "decoded_token": " texture"}, "9436": {"logprob": -3.3846635818481445, "rank": 5, "decoded_token": " setting"}}, {"1046": {"logprob": -0.13945358991622925, "rank": 1, "decoded_token": "."}, "1338": {"logprob": -2.576953649520874, "rank": 2, "decoded_token": ".\n\n"}, "1294": {"logprob": -4.764453411102295, "rank": 3, "decoded_token": " in"}, "1044": {"logprob": -5.701953411102295, "rank": 4, "decoded_token": ","}, "1319": {"logprob": -5.889453411102295, "rank": 5, "decoded_token": " ("}}, {"2": {"logprob": -0.5447223782539368, "rank": 1, "decoded_token": ""}, "1319": {"logprob": -2.607222318649292, "rank": 2, "decoded_token": " ("}, "9380": {"logprob": -3.669722318649292, "rank": 3, "decoded_token": " Here"}, "1531": {"logprob": -3.919722318649292, "rank": 4, "decoded_token": " The"}, "2898": {"logprob": -4.107222557067871, "rank": 5, "decoded_token": " For"}}]]] \ No newline at end of file diff --git a/tests/models/multimodal/generation/test_pixtral.py b/tests/models/multimodal/generation/test_pixtral.py index 48329d9aea3..2ce732342bd 100644 --- a/tests/models/multimodal/generation/test_pixtral.py +++ b/tests/models/multimodal/generation/test_pixtral.py @@ -25,6 +25,7 @@ if TYPE_CHECKING: PIXTRAL_ID = "mistralai/Pixtral-12B-2409" MISTRAL_SMALL_3_1_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503" +MINISTRAL_3B_ID = "mistralai/Ministral-3-3B-Instruct-2512" MODELS = [PIXTRAL_ID, MISTRAL_SMALL_3_1_ID] @@ -116,6 +117,7 @@ assert FIXTURES_PATH.exists() FIXTURE_LOGPROBS_CHAT = { PIXTRAL_ID: FIXTURES_PATH / "pixtral_chat.json", MISTRAL_SMALL_3_1_ID: FIXTURES_PATH / "mistral_small_3_chat.json", + MINISTRAL_3B_ID: FIXTURES_PATH / "ministral_3b_chat.json", } OutputsLogprobs = list[tuple[list[int], str, SampleLogprobs | None]] @@ -209,3 +211,41 @@ def test_chat( name_0="h100_ref", name_1="output", ) + + +@large_gpu_test(min_gb=16) +@pytest.mark.parametrize("dtype", ["bfloat16"]) +def test_chat_consolidated(vllm_runner, dtype: str, local_asset_server) -> None: + EXPECTED_CHAT_LOGPROBS = load_outputs_w_logprobs( + FIXTURE_LOGPROBS_CHAT[MINISTRAL_3B_ID] + ) + with vllm_runner( + MINISTRAL_3B_ID, + dtype=dtype, + tokenizer_mode="mistral", + load_format="mistral", + config_format="mistral", + max_model_len=8192, + limit_mm_per_prompt=LIMIT_MM_PER_PROMPT, + ) as vllm_model: + outputs = [] + urls_all = [local_asset_server.url_for(u) for u in IMG_URLS] + msgs = [ + _create_msg_format(urls_all[:1]), + _create_msg_format(urls_all[:2]), + _create_msg_format(urls_all), + ] + for msg in msgs: + output = vllm_model.llm.chat(msg, sampling_params=SAMPLING_PARAMS) + outputs.extend(output) + + logprobs = vllm_runner._final_steps_generate_w_logprobs(outputs) + for i in range(len(logprobs)): + assert logprobs[i][-1] is None + logprobs[i] = logprobs[i][:-1] + check_logprobs_close( + outputs_0_lst=EXPECTED_CHAT_LOGPROBS, + outputs_1_lst=logprobs, + name_0="h100_ref", + name_1="output", + ) diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index e179638a869..447d6edd986 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -458,13 +458,27 @@ class PixtralForConditionalGeneration( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): _vision_encoder_stacked_params = [ # (param_name, shard_name, shard_id) + # HF format (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), (".gate_up_proj", ".gate_proj", 0), (".gate_up_proj", ".up_proj", 1), + # Mistral native (consolidated) format + (".qkv_proj", ".wq", "q"), + (".qkv_proj", ".wk", "k"), + (".qkv_proj", ".wv", "v"), + (".gate_up_proj", ".w1", 0), + (".gate_up_proj", ".w3", 1), ] + # Remap Mistral native names to HF-style names + # used by the vLLM vision encoder modules. + _vision_encoder_name_remap = { + ".wo.": ".o_proj.", + ".w2.": ".down_proj.", + } + def is_vision_encoder_weights(weight: tuple[str, torch.Tensor]): return weight[0].startswith(("vision_encoder", "vision_tower")) @@ -518,6 +532,11 @@ class PixtralForConditionalGeneration( weight_loader(param, w, shard_id) break else: + for old, new in _vision_encoder_name_remap.items(): + if old in trimmed_name: + trimmed_name = trimmed_name.replace(old, new) + break + param = vision_encoder_dict.get(trimmed_name) if param is not None: weight_loader = getattr( From ec7aafc02ae86c41edf2f245df793a37bd097fc4 Mon Sep 17 00:00:00 2001 From: velonica0 <47554626+velonica0@users.noreply.github.com> Date: Mon, 20 Apr 2026 14:36:59 +0800 Subject: [PATCH 130/150] [CPU][RISC-V] Support multiple RVV VLEN targets via compile-time dispatch (#39478) Signed-off-by: velonica0 --- cmake/cpu_extension.cmake | 49 +- csrc/cpu/cpu_types_riscv.hpp | 839 +--------------------------- csrc/cpu/cpu_types_riscv_defs.hpp | 98 ++++ csrc/cpu/cpu_types_riscv_impl.hpp | 891 ++++++++++++++++++++++++++++++ 4 files changed, 1046 insertions(+), 831 deletions(-) create mode 100644 csrc/cpu/cpu_types_riscv_defs.hpp create mode 100644 csrc/cpu/cpu_types_riscv_impl.hpp diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 698177d0756..8535186cc1e 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -161,16 +161,49 @@ elseif (S390_FOUND) "-mtune=native") elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") message(STATUS "RISC-V detected") - if(RVV_BF16_FOUND) - message(STATUS "BF16 extension detected") - set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d) - add_compile_definitions(RISCV_BF16_SUPPORT) - elseif (RVV_FP16_FOUND) - message(WARNING "BF16 functionality is not available") - set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d) + # VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo + # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256. + if(NOT DEFINED VLLM_RVV_VLEN) + # Auto-detect: find the largest zvlb in /proc/cpuinfo isa line. + if(EXISTS /proc/cpuinfo) + file(READ /proc/cpuinfo _cpuinfo) + set(_best 0) + foreach(_n IN ITEMS 128 256 512 1024) + if(_cpuinfo MATCHES "zvl${_n}b") + set(_best ${_n}) + endif() + endforeach() + if(_best GREATER 0) + set(VLLM_RVV_VLEN ${_best}) + endif() + endif() + # If auto-detect failed (no /proc/cpuinfo or no zvlb reported) + # but the compiler supports RVV, require explicit specification. + if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND)) + message(FATAL_ERROR + "RISC-V RVV is available but VLEN could not be auto-detected. " + "Please specify VLEN explicitly:\n" + " -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n" + " -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)\n" + " -DVLLM_RVV_VLEN=0 (force scalar, no RVV)") + endif() + endif() + if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0) + message(STATUS "RISC-V RVV VLEN=${VLLM_RVV_VLEN}") + if(RVV_BF16_FOUND) + message(STATUS "BF16 extension detected") + set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) + add_compile_definitions(RISCV_BF16_SUPPORT) + elseif(RVV_FP16_FOUND) + message(WARNING "BF16 functionality is not available") + set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) + else() + message(STATUS "compile riscv with scalar (no FP16/BF16)") + set(MARCH_FLAGS -march=rv64gc) + endif() else() message(STATUS "compile riscv with scalar") - list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc") + set(MARCH_FLAGS -march=rv64gc) endif() list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS}) else() diff --git a/csrc/cpu/cpu_types_riscv.hpp b/csrc/cpu/cpu_types_riscv.hpp index 910ee5c1133..e617d98dd00 100644 --- a/csrc/cpu/cpu_types_riscv.hpp +++ b/csrc/cpu/cpu_types_riscv.hpp @@ -1,832 +1,25 @@ #ifndef CPU_TYPES_RISCV_HPP #define CPU_TYPES_RISCV_HPP -#include -#include -#include -#include -#include -#include -#include +// RISC-V Vector (RVV) CPU type definitions for vLLM. +// +// Supports multiple VLENs via compile-time dispatch. The compiler defines +// __riscv_v_min_vlen from the zvlb extension in -march. The defs header +// maps VLEN to the correct LMUL suffixes, and the impl header provides +// VLEN-independent class implementations. +// +// To add support for a new VLEN, add the LMUL mapping in +// cpu_types_riscv_defs.hpp (the impl header needs no changes). -// ============================================================================ -// Vector Register Type Definitions (VLEN=128 bits) -// ============================================================================ - -typedef vfloat16m1_t fixed_vfloat16m1_t - __attribute__((riscv_rvv_vector_bits(128))); -typedef vfloat16m2_t fixed_vfloat16m2_t - __attribute__((riscv_rvv_vector_bits(256))); - -typedef vfloat32m1_t fixed_vfloat32m1_t - __attribute__((riscv_rvv_vector_bits(128))); -typedef vfloat32m2_t fixed_vfloat32m2_t - __attribute__((riscv_rvv_vector_bits(256))); -typedef vfloat32m4_t fixed_vfloat32m4_t - __attribute__((riscv_rvv_vector_bits(512))); -typedef vfloat32m8_t fixed_vfloat32m8_t - __attribute__((riscv_rvv_vector_bits(1024))); - -typedef vint32m2_t fixed_vint32m2_t __attribute__((riscv_rvv_vector_bits(256))); -typedef vint32m4_t fixed_vint32m4_t __attribute__((riscv_rvv_vector_bits(512))); - -typedef vuint16m1_t fixed_vuint16m1_t - __attribute__((riscv_rvv_vector_bits(128))); -typedef vuint16m2_t fixed_vuint16m2_t - __attribute__((riscv_rvv_vector_bits(256))); -typedef vuint16m4_t fixed_vuint16m4_t - __attribute__((riscv_rvv_vector_bits(512))); - -#ifdef RISCV_BF16_SUPPORT -typedef vbfloat16m1_t fixed_vbfloat16m1_t - __attribute__((riscv_rvv_vector_bits(128))); -typedef vbfloat16m2_t fixed_vbfloat16m2_t - __attribute__((riscv_rvv_vector_bits(256))); -typedef vbfloat16m4_t fixed_vbfloat16m4_t - __attribute__((riscv_rvv_vector_bits(512))); +#ifndef __riscv_vector + #error "cpu_types_riscv.hpp included in a non-RVV translation unit" #endif -namespace vec_op { - -#ifdef RISCV_BF16_SUPPORT - #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) -#else - #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) +#ifndef __riscv_v_min_vlen + #error "compiler did not define __riscv_v_min_vlen; pass -march=...zvlb" #endif -#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) +#include "cpu_types_riscv_defs.hpp" +#include "cpu_types_riscv_impl.hpp" -#define FORCE_INLINE __attribute__((always_inline)) inline - -namespace { -template -constexpr void unroll_loop_item(std::integer_sequence, F&& f) { - (f(std::integral_constant{}), ...); -}; -} // namespace - -template >> -constexpr void unroll_loop(F&& f) { - unroll_loop_item(std::make_integer_sequence{}, std::forward(f)); -} - -template -struct Vec { - constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; }; -}; - -struct FP32Vec8; -struct FP32Vec16; - -// ============================================================================ -// FP16 Implementation -// ============================================================================ - -struct FP16Vec8 : public Vec { - constexpr static int VEC_ELEM_NUM = 8; - fixed_vfloat16m1_t reg; - - explicit FP16Vec8(const void* ptr) - : reg(__riscv_vle16_v_f16m1(static_cast(ptr), - VEC_ELEM_NUM)) {}; - - explicit FP16Vec8(const FP32Vec8&); - - void save(void* ptr) const { - __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM); - } - void save(void* ptr, int elem_num) const { - __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, elem_num); - } - void save_strided(void* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(_Float16); - __riscv_vsse16_v_f16m1(static_cast<_Float16*>(ptr), byte_stride, reg, - VEC_ELEM_NUM); - } -}; - -struct FP16Vec16 : public Vec { - constexpr static int VEC_ELEM_NUM = 16; - fixed_vfloat16m2_t reg; - - explicit FP16Vec16(const void* ptr) - : reg(__riscv_vle16_v_f16m2(static_cast(ptr), - VEC_ELEM_NUM)) {}; - - explicit FP16Vec16(const FP32Vec16& vec); - - void save(void* ptr) const { - __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM); - } - void save(void* ptr, int elem_num) const { - __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, elem_num); - } - void save_strided(void* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(_Float16); - __riscv_vsse16_v_f16m2(static_cast<_Float16*>(ptr), byte_stride, reg, - VEC_ELEM_NUM); - } -}; - -// ============================================================================ -// BF16 Implementation -// ============================================================================ - -#ifdef RISCV_BF16_SUPPORT - -FORCE_INLINE fixed_vuint16m1_t bf16_to_u16(fixed_vbfloat16m1_t v) { - return __riscv_vreinterpret_v_bf16m1_u16m1(v); -} -FORCE_INLINE fixed_vuint16m2_t bf16_to_u16(fixed_vbfloat16m2_t v) { - return __riscv_vreinterpret_v_bf16m2_u16m2(v); -} -FORCE_INLINE fixed_vuint16m4_t bf16_to_u16(fixed_vbfloat16m4_t v) { - return __riscv_vreinterpret_v_bf16m4_u16m4(v); -} - -struct BF16Vec8 : public Vec { - constexpr static int VEC_ELEM_NUM = 8; - fixed_vbfloat16m1_t reg; - - explicit BF16Vec8(const void* ptr) - : reg(__riscv_vreinterpret_v_u16m1_bf16m1(__riscv_vle16_v_u16m1( - reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; - - explicit BF16Vec8(fixed_vbfloat16m1_t data) : reg(data) {}; - explicit BF16Vec8(const FP32Vec8&); - - void save(void* ptr) const { - __riscv_vse16_v_u16m1(reinterpret_cast(ptr), bf16_to_u16(reg), - VEC_ELEM_NUM); - } - void save(void* ptr, int elem_num) const { - __riscv_vse16_v_u16m1(reinterpret_cast(ptr), bf16_to_u16(reg), - elem_num); - } - void save_strided(void* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - __riscv_vsse16_v_u16m1(reinterpret_cast(ptr), byte_stride, - bf16_to_u16(reg), VEC_ELEM_NUM); - } -}; - -struct BF16Vec16 : public Vec { - constexpr static int VEC_ELEM_NUM = 16; - fixed_vbfloat16m2_t reg; - - explicit BF16Vec16(const void* ptr) - : reg(__riscv_vreinterpret_v_u16m2_bf16m2(__riscv_vle16_v_u16m2( - reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; - - explicit BF16Vec16(fixed_vbfloat16m2_t data) : reg(data) {}; - explicit BF16Vec16(const FP32Vec16&); - - void save(void* ptr) const { - __riscv_vse16_v_u16m2(reinterpret_cast(ptr), bf16_to_u16(reg), - VEC_ELEM_NUM); - } - void save(void* ptr, int elem_num) const { - __riscv_vse16_v_u16m2(reinterpret_cast(ptr), bf16_to_u16(reg), - elem_num); - } - void save_strided(void* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - __riscv_vsse16_v_u16m2(reinterpret_cast(ptr), byte_stride, - bf16_to_u16(reg), VEC_ELEM_NUM); - } -}; - -struct BF16Vec32 : public Vec { - constexpr static int VEC_ELEM_NUM = 32; - fixed_vbfloat16m4_t reg; - - explicit BF16Vec32(const void* ptr) - : reg(__riscv_vreinterpret_v_u16m4_bf16m4(__riscv_vle16_v_u16m4( - reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; - - explicit BF16Vec32(fixed_vbfloat16m4_t data) : reg(data) {}; - - explicit BF16Vec32(const BF16Vec8& v) { - fixed_vuint16m1_t u16_val = bf16_to_u16(v.reg); - fixed_vuint16m4_t u16_combined = - __riscv_vcreate_v_u16m1_u16m4(u16_val, u16_val, u16_val, u16_val); - reg = __riscv_vreinterpret_v_u16m4_bf16m4(u16_combined); - }; - - void save(void* ptr) const { - __riscv_vse16_v_u16m4(reinterpret_cast(ptr), bf16_to_u16(reg), - VEC_ELEM_NUM); - } - void save(void* ptr, int elem_num) const { - __riscv_vse16_v_u16m4(reinterpret_cast(ptr), bf16_to_u16(reg), - elem_num); - } - void save_strided(void* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - __riscv_vsse16_v_u16m4(reinterpret_cast(ptr), byte_stride, - bf16_to_u16(reg), VEC_ELEM_NUM); - } -}; - -#else -// ============================================================================ -// BF16 Fallback Implementation (FP32 Simulation) -// ============================================================================ - -struct BF16Vec8 : public Vec { - constexpr static int VEC_ELEM_NUM = 8; - fixed_vfloat32m2_t reg_fp32; - explicit BF16Vec8(const void* ptr) { - const uint16_t* u16 = static_cast(ptr); - float tmp[8]; - for (int i = 0; i < 8; ++i) { - uint32_t v = static_cast(u16[i]) << 16; - std::memcpy(&tmp[i], &v, 4); - } - reg_fp32 = __riscv_vle32_v_f32m2(tmp, 8); - } - explicit BF16Vec8(const FP32Vec8&); - void save(void* ptr) const { - float tmp[8]; - __riscv_vse32_v_f32m2(tmp, reg_fp32, 8); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < 8; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - void save(void* ptr, int elem_num) const { - float tmp[8]; - __riscv_vse32_v_f32m2(tmp, reg_fp32, 8); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < elem_num; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - void save_strided(void* ptr, ptrdiff_t stride) const { - float tmp[8]; - __riscv_vse32_v_f32m2(tmp, reg_fp32, 8); - uint8_t* u8 = static_cast(ptr); - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - for (int i = 0; i < 8; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - uint16_t val = static_cast(v >> 16); - *reinterpret_cast(u8 + i * byte_stride) = val; - } - } -}; - -struct BF16Vec16 : public Vec { - constexpr static int VEC_ELEM_NUM = 16; - fixed_vfloat32m4_t reg_fp32; - explicit BF16Vec16(const void* ptr) { - const uint16_t* u16 = static_cast(ptr); - float tmp[16]; - for (int i = 0; i < 16; ++i) { - uint32_t v = static_cast(u16[i]) << 16; - std::memcpy(&tmp[i], &v, 4); - } - reg_fp32 = __riscv_vle32_v_f32m4(tmp, 16); - } - explicit BF16Vec16(const FP32Vec16&); - void save(void* ptr) const { - float tmp[16]; - __riscv_vse32_v_f32m4(tmp, reg_fp32, 16); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < 16; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - void save(void* ptr, int elem_num) const { - float tmp[16]; - __riscv_vse32_v_f32m4(tmp, reg_fp32, 16); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < elem_num; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - void save_strided(void* ptr, ptrdiff_t stride) const { - float tmp[16]; - __riscv_vse32_v_f32m4(tmp, reg_fp32, 16); - uint8_t* u8 = static_cast(ptr); - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - for (int i = 0; i < 16; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - uint16_t val = static_cast(v >> 16); - *reinterpret_cast(u8 + i * byte_stride) = val; - } - } -}; - -struct BF16Vec32 : public Vec { - constexpr static int VEC_ELEM_NUM = 32; - fixed_vfloat32m8_t reg_fp32; - - explicit BF16Vec32(const void* ptr) { - const uint16_t* u16 = static_cast(ptr); - float tmp[32]; - for (int i = 0; i < 32; ++i) { - uint32_t v = static_cast(u16[i]) << 16; - std::memcpy(&tmp[i], &v, 4); - } - reg_fp32 = __riscv_vle32_v_f32m8(tmp, 32); - } - - explicit BF16Vec32(const BF16Vec8& v) { - float tmp_small[8]; - __riscv_vse32_v_f32m2(tmp_small, v.reg_fp32, 8); - float tmp_large[32]; - for (int i = 0; i < 4; ++i) { - std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float)); - } - reg_fp32 = __riscv_vle32_v_f32m8(tmp_large, 32); - } - - void save(void* ptr) const { - float tmp[32]; - __riscv_vse32_v_f32m8(tmp, reg_fp32, 32); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < 32; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - - void save(void* ptr, int elem_num) const { - float tmp[32]; - __riscv_vse32_v_f32m8(tmp, reg_fp32, 32); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < elem_num; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - - void save_strided(void* ptr, ptrdiff_t stride) const { - float tmp[32]; - __riscv_vse32_v_f32m8(tmp, reg_fp32, 32); - uint8_t* u8 = static_cast(ptr); - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - for (int i = 0; i < 32; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - uint16_t val = static_cast(v >> 16); - *reinterpret_cast(u8 + i * byte_stride) = val; - } - } -}; -#endif - -// ============================================================================ -// FP32 Implementation -// ============================================================================ - -struct FP32Vec4 : public Vec { - constexpr static int VEC_ELEM_NUM = 4; - fixed_vfloat32m1_t reg; - explicit FP32Vec4(float v) : reg(__riscv_vfmv_v_f_f32m1(v, VEC_ELEM_NUM)) {}; - explicit FP32Vec4() : reg(__riscv_vfmv_v_f_f32m1(0.0f, VEC_ELEM_NUM)) {}; - explicit FP32Vec4(const float* ptr) - : reg(__riscv_vle32_v_f32m1(ptr, VEC_ELEM_NUM)) {}; - explicit FP32Vec4(fixed_vfloat32m1_t data) : reg(data) {}; - explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}; - void save(float* ptr) const { __riscv_vse32_v_f32m1(ptr, reg, VEC_ELEM_NUM); } - void save(float* ptr, int elem_num) const { - __riscv_vse32_v_f32m1(ptr, reg, elem_num); - } -}; - -struct FP32Vec8 : public Vec { - constexpr static int VEC_ELEM_NUM = 8; - fixed_vfloat32m2_t reg; - - explicit FP32Vec8(float v) : reg(__riscv_vfmv_v_f_f32m2(v, VEC_ELEM_NUM)) {}; - explicit FP32Vec8() : reg(__riscv_vfmv_v_f_f32m2(0.0f, VEC_ELEM_NUM)) {}; - explicit FP32Vec8(const float* ptr) - : reg(__riscv_vle32_v_f32m2(ptr, VEC_ELEM_NUM)) {}; - explicit FP32Vec8(fixed_vfloat32m2_t data) : reg(data) {}; - explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {}; - explicit FP32Vec8(const FP16Vec8& v) - : reg(__riscv_vfwcvt_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {}; - explicit FP32Vec8(fixed_vfloat16m1_t v) - : reg(__riscv_vfwcvt_f_f_v_f32m2(v, VEC_ELEM_NUM)) {}; - -#ifdef RISCV_BF16_SUPPORT - explicit FP32Vec8(fixed_vbfloat16m1_t v) - : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v, VEC_ELEM_NUM)) {}; - explicit FP32Vec8(const BF16Vec8& v) - : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {}; -#else - explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {}; -#endif - - float reduce_sum() const { - fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); - scalar = __riscv_vfredusum_vs_f32m2_f32m1(reg, scalar, VEC_ELEM_NUM); - return __riscv_vfmv_f_s_f32m1_f32(scalar); - } - - FP32Vec8 operator*(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfmul_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec8 operator+(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfadd_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec8 operator-(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfsub_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec8 operator/(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfdiv_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - - FP32Vec8 min(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec8 max(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec8 abs() const { - return FP32Vec8(__riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM)); - } - - FP32Vec8 min(const FP32Vec8& b, int elem_num) const { - return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, elem_num)); - } - FP32Vec8 max(const FP32Vec8& b, int elem_num) const { - return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, elem_num)); - } - - FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const { - fixed_vfloat32m2_t temp = - __riscv_vfmax_vv_f32m2(min_v.reg, reg, VEC_ELEM_NUM); - return FP32Vec8(__riscv_vfmin_vv_f32m2(max_v.reg, temp, VEC_ELEM_NUM)); - } - - void save(float* ptr) const { __riscv_vse32_v_f32m2(ptr, reg, VEC_ELEM_NUM); } - void save(float* ptr, int elem_num) const { - __riscv_vse32_v_f32m2(ptr, reg, elem_num); - } - void save_strided(float* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(float); - __riscv_vsse32_v_f32m2(ptr, byte_stride, reg, VEC_ELEM_NUM); - } - - FP32Vec8 exp() const { - const float inv_ln2 = 1.44269504088896341f; - fixed_vfloat32m2_t x_scaled = - __riscv_vfmul_vf_f32m2(reg, inv_ln2, VEC_ELEM_NUM); - fixed_vint32m2_t n_int = __riscv_vfcvt_x_f_v_i32m2(x_scaled, VEC_ELEM_NUM); - fixed_vfloat32m2_t n_float = __riscv_vfcvt_f_x_v_f32m2(n_int, VEC_ELEM_NUM); - - fixed_vfloat32m2_t r = - __riscv_vfsub_vv_f32m2(x_scaled, n_float, VEC_ELEM_NUM); - - fixed_vfloat32m2_t poly = - __riscv_vfmv_v_f_f32m2(0.001333355810164f, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(poly, 0.009618129107628f, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(poly, 0.055504108664821f, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(poly, 0.240226506959101f, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(poly, 0.693147180559945f, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(poly, 1.0f, VEC_ELEM_NUM); - - fixed_vint32m2_t biased_exp = - __riscv_vadd_vx_i32m2(n_int, 127, VEC_ELEM_NUM); - biased_exp = __riscv_vmax_vx_i32m2(biased_exp, 0, VEC_ELEM_NUM); - fixed_vint32m2_t exponent_bits = - __riscv_vsll_vx_i32m2(biased_exp, 23, VEC_ELEM_NUM); - fixed_vfloat32m2_t scale = - __riscv_vreinterpret_v_i32m2_f32m2(exponent_bits); - - return FP32Vec8(__riscv_vfmul_vv_f32m2(poly, scale, VEC_ELEM_NUM)); - } - - FP32Vec8 tanh() const { - fixed_vfloat32m2_t x_clamped = __riscv_vfmin_vf_f32m2( - __riscv_vfmax_vf_f32m2(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM); - fixed_vfloat32m2_t x2 = - __riscv_vfmul_vf_f32m2(x_clamped, 2.0f, VEC_ELEM_NUM); - FP32Vec8 exp_val = FP32Vec8(x2).exp(); - fixed_vfloat32m2_t num = - __riscv_vfsub_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM); - fixed_vfloat32m2_t den = - __riscv_vfadd_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM); - return FP32Vec8(__riscv_vfdiv_vv_f32m2(num, den, VEC_ELEM_NUM)); - } - - FP32Vec8 er() const { - const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f, - a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f; - fixed_vfloat32m2_t abs_x = __riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM); - - fixed_vfloat32m2_t t = __riscv_vfadd_vf_f32m2( - __riscv_vfmul_vf_f32m2(abs_x, p, VEC_ELEM_NUM), 1.0f, VEC_ELEM_NUM); - t = __riscv_vfrdiv_vf_f32m2(t, 1.0f, VEC_ELEM_NUM); - - fixed_vfloat32m2_t poly = __riscv_vfmv_v_f_f32m2(a5, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM), - a4, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM), - a3, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM), - a2, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM), - a1, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM); - - fixed_vfloat32m2_t exp_val = - FP32Vec8(__riscv_vfneg_v_f32m2( - __riscv_vfmul_vv_f32m2(abs_x, abs_x, VEC_ELEM_NUM), - VEC_ELEM_NUM)) - .exp() - .reg; - fixed_vfloat32m2_t res = __riscv_vfrsub_vf_f32m2( - __riscv_vfmul_vv_f32m2(poly, exp_val, VEC_ELEM_NUM), 1.0f, - VEC_ELEM_NUM); - - vbool16_t mask = __riscv_vmflt_vf_f32m2_b16(reg, 0.0f, VEC_ELEM_NUM); - return FP32Vec8(__riscv_vfneg_v_f32m2_m(mask, res, VEC_ELEM_NUM)); - } -}; - -struct FP32Vec16 : public Vec { - constexpr static int VEC_ELEM_NUM = 16; - fixed_vfloat32m4_t reg; - - explicit FP32Vec16(float v) : reg(__riscv_vfmv_v_f_f32m4(v, VEC_ELEM_NUM)) {}; - explicit FP32Vec16() : reg(__riscv_vfmv_v_f_f32m4(0.0f, VEC_ELEM_NUM)) {}; - explicit FP32Vec16(const float* ptr) - : reg(__riscv_vle32_v_f32m4(ptr, VEC_ELEM_NUM)) {}; - explicit FP32Vec16(fixed_vfloat32m4_t data) : reg(data) {}; - explicit FP32Vec16(const FP32Vec8& data) - : reg(__riscv_vcreate_v_f32m2_f32m4(data.reg, data.reg)) {}; - explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {}; - explicit FP32Vec16(const FP16Vec16& v); - -#ifdef RISCV_BF16_SUPPORT - explicit FP32Vec16(fixed_vbfloat16m2_t v) - : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v, VEC_ELEM_NUM)) {}; - explicit FP32Vec16(const BF16Vec16& v) - : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v.reg, VEC_ELEM_NUM)) {}; -#else - explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {}; -#endif - - FP32Vec16 operator+(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfadd_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec16 operator-(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfsub_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec16 operator*(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfmul_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec16 operator/(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfdiv_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - - FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfmacc_vv_f32m4(reg, a.reg, b.reg, VEC_ELEM_NUM)); - } - - float reduce_sum() const { - fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); - scalar = __riscv_vfredusum_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM); - return __riscv_vfmv_f_s_f32m1_f32(scalar); - } - - float reduce_max() const { - fixed_vfloat32m1_t scalar = - __riscv_vfmv_s_f_f32m1(std::numeric_limits::lowest(), 1); - scalar = __riscv_vfredmax_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM); - return __riscv_vfmv_f_s_f32m1_f32(scalar); - } - - float reduce_min() const { - fixed_vfloat32m1_t scalar = - __riscv_vfmv_s_f_f32m1(std::numeric_limits::max(), 1); - scalar = __riscv_vfredmin_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM); - return __riscv_vfmv_f_s_f32m1_f32(scalar); - } - - template - float reduce_sub_sum(int idx) { - static_assert(VEC_ELEM_NUM % group_size == 0); - const int start = idx * group_size; - vuint32m4_t indices = __riscv_vid_v_u32m4(VEC_ELEM_NUM); - vbool8_t mask = __riscv_vmand_mm_b8( - __riscv_vmsgeu_vx_u32m4_b8(indices, start, VEC_ELEM_NUM), - __riscv_vmsltu_vx_u32m4_b8(indices, start + group_size, VEC_ELEM_NUM), - VEC_ELEM_NUM); - fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); - scalar = - __riscv_vfredusum_vs_f32m4_f32m1_m(mask, reg, scalar, VEC_ELEM_NUM); - return __riscv_vfmv_f_s_f32m1_f32(scalar); - }; - - FP32Vec16 max(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfmax_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec16 min(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfmin_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec16 abs() const { - return FP32Vec16(__riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM)); - } - - FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const { - return FP32Vec16(__riscv_vfmin_vv_f32m4( - max_v.reg, __riscv_vfmax_vv_f32m4(min_v.reg, reg, VEC_ELEM_NUM), - VEC_ELEM_NUM)); - } - - void save(float* ptr) const { __riscv_vse32_v_f32m4(ptr, reg, VEC_ELEM_NUM); } - void save(float* ptr, int elem_num) const { - __riscv_vse32_v_f32m4(ptr, reg, elem_num); - } - void save_strided(float* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(float); - __riscv_vsse32_v_f32m4(ptr, byte_stride, reg, VEC_ELEM_NUM); - } - - FP32Vec16 exp() const { - const float inv_ln2 = 1.44269504088896341f; - fixed_vfloat32m4_t x_scaled = - __riscv_vfmul_vf_f32m4(reg, inv_ln2, VEC_ELEM_NUM); - fixed_vint32m4_t n_int = __riscv_vfcvt_x_f_v_i32m4(x_scaled, VEC_ELEM_NUM); - fixed_vfloat32m4_t n_float = __riscv_vfcvt_f_x_v_f32m4(n_int, VEC_ELEM_NUM); - fixed_vfloat32m4_t r = - __riscv_vfsub_vv_f32m4(x_scaled, n_float, VEC_ELEM_NUM); - - fixed_vfloat32m4_t poly = - __riscv_vfmv_v_f_f32m4(0.001333355810164f, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM), - 0.009618129107628f, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM), - 0.055504108664821f, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM), - 0.240226506959101f, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM), - 0.693147180559945f, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM), - 1.0f, VEC_ELEM_NUM); - - fixed_vint32m4_t biased_exp = __riscv_vmax_vx_i32m4( - __riscv_vadd_vx_i32m4(n_int, 127, VEC_ELEM_NUM), 0, VEC_ELEM_NUM); - fixed_vfloat32m4_t scale = __riscv_vreinterpret_v_i32m4_f32m4( - __riscv_vsll_vx_i32m4(biased_exp, 23, VEC_ELEM_NUM)); - - return FP32Vec16(__riscv_vfmul_vv_f32m4(poly, scale, VEC_ELEM_NUM)); - } - - FP32Vec16 tanh() const { - fixed_vfloat32m4_t x_clamped = __riscv_vfmin_vf_f32m4( - __riscv_vfmax_vf_f32m4(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM); - FP32Vec16 exp_val = - FP32Vec16(__riscv_vfmul_vf_f32m4(x_clamped, 2.0f, VEC_ELEM_NUM)).exp(); - return FP32Vec16(__riscv_vfdiv_vv_f32m4( - __riscv_vfsub_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM), - __riscv_vfadd_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM), VEC_ELEM_NUM)); - } - - FP32Vec16 er() const { - const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f, - a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f; - fixed_vfloat32m4_t abs_x = __riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM); - fixed_vfloat32m4_t t = __riscv_vfrdiv_vf_f32m4( - __riscv_vfadd_vf_f32m4(__riscv_vfmul_vf_f32m4(abs_x, p, VEC_ELEM_NUM), - 1.0f, VEC_ELEM_NUM), - 1.0f, VEC_ELEM_NUM); - - fixed_vfloat32m4_t poly = __riscv_vfmv_v_f_f32m4(a5, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM), - a4, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM), - a3, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM), - a2, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM), - a1, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM); - - fixed_vfloat32m4_t exp_val = - FP32Vec16(__riscv_vfneg_v_f32m4( - __riscv_vfmul_vv_f32m4(abs_x, abs_x, VEC_ELEM_NUM), - VEC_ELEM_NUM)) - .exp() - .reg; - fixed_vfloat32m4_t res = __riscv_vfrsub_vf_f32m4( - __riscv_vfmul_vv_f32m4(poly, exp_val, VEC_ELEM_NUM), 1.0f, - VEC_ELEM_NUM); - - vbool8_t mask = __riscv_vmflt_vf_f32m4_b8(reg, 0.0f, VEC_ELEM_NUM); - return FP32Vec16(__riscv_vfneg_v_f32m4_m(mask, res, VEC_ELEM_NUM)); - } -}; - -// ============================================================================ -// Type Traits & Global Helpers -// ============================================================================ - -template -struct VecType { - using vec_type = void; - using vec_t = void; -}; - -template -using vec_t = typename VecType::vec_type; - -template <> -struct VecType { - using vec_type = FP32Vec8; - using vec_t = FP32Vec8; -}; -template <> -struct VecType { - using vec_type = FP16Vec8; - using vec_t = FP16Vec8; -}; -template <> -struct VecType { - using vec_type = BF16Vec8; - using vec_t = BF16Vec8; -}; - -template -void storeFP32(float v, T* ptr) { - *ptr = v; -} -template <> -inline void storeFP32(float v, c10::Half* ptr) { - *reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v); -} - -inline FP16Vec16::FP16Vec16(const FP32Vec16& v) { - reg = __riscv_vfncvt_f_f_w_f16m2(v.reg, VEC_ELEM_NUM); -} -inline FP16Vec8::FP16Vec8(const FP32Vec8& v) { - reg = __riscv_vfncvt_f_f_w_f16m1(v.reg, VEC_ELEM_NUM); -} -inline FP32Vec16::FP32Vec16(const FP16Vec16& v) { - reg = __riscv_vfwcvt_f_f_v_f32m4(v.reg, VEC_ELEM_NUM); -} -inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) { - acc = acc.fma(a, b); -} - -#ifdef RISCV_BF16_SUPPORT -template <> -inline void storeFP32(float v, c10::BFloat16* ptr) { - *ptr = static_cast<__bf16>(v); -}; -inline BF16Vec8::BF16Vec8(const FP32Vec8& v) - : reg(__riscv_vfncvtbf16_f_f_w_bf16m1(v.reg, VEC_ELEM_NUM)) {}; -inline BF16Vec16::BF16Vec16(const FP32Vec16& v) - : reg(__riscv_vfncvtbf16_f_f_w_bf16m2(v.reg, VEC_ELEM_NUM)) {}; -#else -template <> -inline void storeFP32(float v, c10::BFloat16* ptr) { - uint32_t val; - std::memcpy(&val, &v, 4); - *reinterpret_cast(ptr) = static_cast(val >> 16); -} -inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {} -inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {} -#endif - -inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); } - -} // namespace vec_op - -#ifndef CPU_KERNEL_GUARD_IN - #define CPU_KERNEL_GUARD_IN(NAME) -#endif - -#ifndef CPU_KERNEL_GUARD_OUT - #define CPU_KERNEL_GUARD_OUT(NAME) -#endif - -#endif // CPU_TYPES_RISCV_HPP \ No newline at end of file +#endif // CPU_TYPES_RISCV_HPP diff --git a/csrc/cpu/cpu_types_riscv_defs.hpp b/csrc/cpu/cpu_types_riscv_defs.hpp new file mode 100644 index 00000000000..04ce2728f03 --- /dev/null +++ b/csrc/cpu/cpu_types_riscv_defs.hpp @@ -0,0 +1,98 @@ +#ifndef CPU_TYPES_RISCV_DEFS_HPP +#define CPU_TYPES_RISCV_DEFS_HPP + +// VLEN-to-LMUL mapping for RISC-V Vector extension. +// +// LMUL_ expands to the LMUL suffix giving N total bits of vector data: +// VLEN=128: LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8 +// VLEN=256: LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4 + +#include + +#if __riscv_v_min_vlen == 128 + #define LMUL_128 m1 + #define LMUL_256 m2 + #define LMUL_512 m4 + #define LMUL_1024 m8 + #define BOOL_256 b16 + #define BOOL_512 b8 +#elif __riscv_v_min_vlen == 256 + #define LMUL_128 mf2 + #define LMUL_256 m1 + #define LMUL_512 m2 + #define LMUL_1024 m4 + #define BOOL_256 b32 + #define BOOL_512 b16 +#else + #error "cpu_types_riscv_defs.hpp: unsupported __riscv_v_min_vlen" +#endif + +// Token-paste helpers. +#define _RVV_P2(a, b) a##b +#define _RVV_P3(a, b, c) a##b##c +#define _RVV_P4(a, b, c, d) a##b##c##d +#define RVVTYPE(base, lmul, suffix) _RVV_P3(base, lmul, suffix) +#define RVVI(base, lmul) _RVV_P2(base, lmul) +#define RVVI3(base, lmul, suffix) _RVV_P3(base, lmul, suffix) +#define RVVI4(a, b, c, d) _RVV_P4(a, b, c, d) +// For mask intrinsics: RVVIB(base, LMUL_256, BOOL_256) → base##m2##_##b16 +#define _RVV_PB(base, lmul, btype) base##lmul##_##btype +#define RVVIB(base, lmul, btype) _RVV_PB(base, lmul, btype) + +// ---- Semantic fixed-vector typedefs (named by element count) ---- + +// float16 +typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t + __attribute__((riscv_rvv_vector_bits(128))); +typedef RVVTYPE(vfloat16, LMUL_256, _t) fixed_fp16x16_t + __attribute__((riscv_rvv_vector_bits(256))); + +// float32 +typedef RVVTYPE(vfloat32, LMUL_128, _t) fixed_fp32x4_t + __attribute__((riscv_rvv_vector_bits(128))); +typedef RVVTYPE(vfloat32, LMUL_256, _t) fixed_fp32x8_t + __attribute__((riscv_rvv_vector_bits(256))); +typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t + __attribute__((riscv_rvv_vector_bits(512))); +typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t + __attribute__((riscv_rvv_vector_bits(1024))); + +// int32 +typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t + __attribute__((riscv_rvv_vector_bits(256))); +typedef RVVTYPE(vint32, LMUL_512, _t) fixed_i32x16_t + __attribute__((riscv_rvv_vector_bits(512))); + +// uint16 +typedef RVVTYPE(vuint16, LMUL_128, _t) fixed_u16x8_t + __attribute__((riscv_rvv_vector_bits(128))); +typedef RVVTYPE(vuint16, LMUL_256, _t) fixed_u16x16_t + __attribute__((riscv_rvv_vector_bits(256))); +typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t + __attribute__((riscv_rvv_vector_bits(512))); + +// bfloat16 +#ifdef RISCV_BF16_SUPPORT +typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t + __attribute__((riscv_rvv_vector_bits(128))); +typedef RVVTYPE(vbfloat16, LMUL_256, _t) fixed_bf16x16_t + __attribute__((riscv_rvv_vector_bits(256))); +typedef RVVTYPE(vbfloat16, LMUL_512, _t) fixed_bf16x32_t + __attribute__((riscv_rvv_vector_bits(512))); +#endif + +// ---- Reduction accumulator type (always m1 = one register of f32) ---- +// Used for scalar reductions; only element [0] is meaningful. +typedef vfloat32m1_t rvv_f32_accum_t + __attribute__((riscv_rvv_vector_bits(__riscv_v_min_vlen))); + +// ---- Mask types for f32 elements ---- +#if __riscv_v_min_vlen == 128 +typedef vbool16_t rvv_mask_f32x8_t; +typedef vbool8_t rvv_mask_f32x16_t; +#elif __riscv_v_min_vlen == 256 +typedef vbool32_t rvv_mask_f32x8_t; +typedef vbool16_t rvv_mask_f32x16_t; +#endif + +#endif // CPU_TYPES_RISCV_DEFS_HPP diff --git a/csrc/cpu/cpu_types_riscv_impl.hpp b/csrc/cpu/cpu_types_riscv_impl.hpp new file mode 100644 index 00000000000..7c25ccc5059 --- /dev/null +++ b/csrc/cpu/cpu_types_riscv_impl.hpp @@ -0,0 +1,891 @@ +#ifndef CPU_TYPES_RISCV_IMPL_HPP +#define CPU_TYPES_RISCV_IMPL_HPP + +// Shared implementation of RVV vector-type wrapper classes. +// This file is VLEN-independent: it uses the semantic type names and +// RVVI() intrinsic macros from cpu_types_riscv_defs.hpp. +// +// DO NOT include this file directly; include cpu_types_riscv.hpp instead. + +#include +#include +#include +#include +#include +#include +namespace vec_op { + +#ifdef RISCV_BF16_SUPPORT + #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +#else + #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) +#endif + +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +#define FORCE_INLINE __attribute__((always_inline)) inline + +namespace { +template +constexpr void unroll_loop_item(std::integer_sequence, F&& f) { + (f(std::integral_constant{}), ...); +}; +} // namespace + +template >> +constexpr void unroll_loop(F&& f) { + unroll_loop_item(std::make_integer_sequence{}, std::forward(f)); +} + +template +struct Vec { + constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; }; +}; + +struct FP32Vec8; +struct FP32Vec16; + +// ============================================================================ +// FP16 Implementation +// ============================================================================ + +struct FP16Vec8 : public Vec { + constexpr static int VEC_ELEM_NUM = 8; + fixed_fp16x8_t reg; + + explicit FP16Vec8(const void* ptr) + : reg(RVVI(__riscv_vle16_v_f16, LMUL_128)( + static_cast(ptr), VEC_ELEM_NUM)) {}; + + explicit FP16Vec8(const FP32Vec8&); + + void save(void* ptr) const { + RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg, + VEC_ELEM_NUM); + } + void save(void* ptr, int elem_num) const { + RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg, + elem_num); + } + void save_strided(void* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(_Float16); + RVVI(__riscv_vsse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), + byte_stride, reg, VEC_ELEM_NUM); + } +}; + +struct FP16Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + fixed_fp16x16_t reg; + + explicit FP16Vec16(const void* ptr) + : reg(RVVI(__riscv_vle16_v_f16, LMUL_256)( + static_cast(ptr), VEC_ELEM_NUM)) {}; + + explicit FP16Vec16(const FP32Vec16& vec); + + void save(void* ptr) const { + RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg, + VEC_ELEM_NUM); + } + void save(void* ptr, int elem_num) const { + RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg, + elem_num); + } + void save_strided(void* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(_Float16); + RVVI(__riscv_vsse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), + byte_stride, reg, VEC_ELEM_NUM); + } +}; + +// ============================================================================ +// BF16 Implementation +// ============================================================================ + +#ifdef RISCV_BF16_SUPPORT + +FORCE_INLINE fixed_u16x8_t bf16_to_u16(fixed_bf16x8_t v) { + return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_128, _u16, LMUL_128)(v); +} +FORCE_INLINE fixed_u16x16_t bf16_to_u16(fixed_bf16x16_t v) { + return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_256, _u16, LMUL_256)(v); +} +FORCE_INLINE fixed_u16x32_t bf16_to_u16(fixed_bf16x32_t v) { + return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_512, _u16, LMUL_512)(v); +} + +struct BF16Vec8 : public Vec { + constexpr static int VEC_ELEM_NUM = 8; + fixed_bf16x8_t reg; + + explicit BF16Vec8(const void* ptr) + : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16, + LMUL_128)(RVVI(__riscv_vle16_v_u16, LMUL_128)( + reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; + + explicit BF16Vec8(fixed_bf16x8_t data) : reg(data) {}; + explicit BF16Vec8(const FP32Vec8&); + + void save(void* ptr) const { + RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast(ptr), + bf16_to_u16(reg), VEC_ELEM_NUM); + } + void save(void* ptr, int elem_num) const { + RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast(ptr), + bf16_to_u16(reg), elem_num); + } + void save_strided(void* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + RVVI(__riscv_vsse16_v_u16, LMUL_128)(reinterpret_cast(ptr), + byte_stride, bf16_to_u16(reg), + VEC_ELEM_NUM); + } +}; + +struct BF16Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + fixed_bf16x16_t reg; + + explicit BF16Vec16(const void* ptr) + : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16, + LMUL_256)(RVVI(__riscv_vle16_v_u16, LMUL_256)( + reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; + + explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {}; + explicit BF16Vec16(const FP32Vec16&); + + void save(void* ptr) const { + RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast(ptr), + bf16_to_u16(reg), VEC_ELEM_NUM); + } + void save(void* ptr, int elem_num) const { + RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast(ptr), + bf16_to_u16(reg), elem_num); + } + void save_strided(void* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + RVVI(__riscv_vsse16_v_u16, LMUL_256)(reinterpret_cast(ptr), + byte_stride, bf16_to_u16(reg), + VEC_ELEM_NUM); + } +}; + +struct BF16Vec32 : public Vec { + constexpr static int VEC_ELEM_NUM = 32; + fixed_bf16x32_t reg; + + explicit BF16Vec32(const void* ptr) + : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, + LMUL_512)(RVVI(__riscv_vle16_v_u16, LMUL_512)( + reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; + + explicit BF16Vec32(fixed_bf16x32_t data) : reg(data) {}; + + explicit BF16Vec32(const BF16Vec8& v) { + fixed_u16x8_t u16_val = bf16_to_u16(v.reg); + fixed_u16x32_t u16_combined = + RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)( + u16_val, u16_val, u16_val, u16_val); + reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, + LMUL_512)(u16_combined); + }; + + void save(void* ptr) const { + RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast(ptr), + bf16_to_u16(reg), VEC_ELEM_NUM); + } + void save(void* ptr, int elem_num) const { + RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast(ptr), + bf16_to_u16(reg), elem_num); + } + void save_strided(void* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + RVVI(__riscv_vsse16_v_u16, LMUL_512)(reinterpret_cast(ptr), + byte_stride, bf16_to_u16(reg), + VEC_ELEM_NUM); + } +}; + +#else +// ============================================================================ +// BF16 Fallback Implementation (FP32 Simulation) +// ============================================================================ + +struct BF16Vec8 : public Vec { + constexpr static int VEC_ELEM_NUM = 8; + fixed_fp32x8_t reg_fp32; + explicit BF16Vec8(const void* ptr) { + const uint16_t* u16 = static_cast(ptr); + float tmp[8]; + for (int i = 0; i < 8; ++i) { + uint32_t v = static_cast(u16[i]) << 16; + std::memcpy(&tmp[i], &v, 4); + } + reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8); + } + explicit BF16Vec8(const FP32Vec8&); + void save(void* ptr) const { + float tmp[8]; + RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < 8; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + void save(void* ptr, int elem_num) const { + float tmp[8]; + RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < elem_num; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + void save_strided(void* ptr, ptrdiff_t stride) const { + float tmp[8]; + RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8); + uint8_t* u8 = static_cast(ptr); + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + for (int i = 0; i < 8; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + uint16_t val = static_cast(v >> 16); + *reinterpret_cast(u8 + i * byte_stride) = val; + } + } +}; + +struct BF16Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + fixed_fp32x16_t reg_fp32; + explicit BF16Vec16(const void* ptr) { + const uint16_t* u16 = static_cast(ptr); + float tmp[16]; + for (int i = 0; i < 16; ++i) { + uint32_t v = static_cast(u16[i]) << 16; + std::memcpy(&tmp[i], &v, 4); + } + reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16); + } + explicit BF16Vec16(const FP32Vec16&); + void save(void* ptr) const { + float tmp[16]; + RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < 16; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + void save(void* ptr, int elem_num) const { + float tmp[16]; + RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < elem_num; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + void save_strided(void* ptr, ptrdiff_t stride) const { + float tmp[16]; + RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16); + uint8_t* u8 = static_cast(ptr); + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + for (int i = 0; i < 16; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + uint16_t val = static_cast(v >> 16); + *reinterpret_cast(u8 + i * byte_stride) = val; + } + } +}; + +struct BF16Vec32 : public Vec { + constexpr static int VEC_ELEM_NUM = 32; + fixed_fp32x32_t reg_fp32; + + explicit BF16Vec32(const void* ptr) { + const uint16_t* u16 = static_cast(ptr); + float tmp[32]; + for (int i = 0; i < 32; ++i) { + uint32_t v = static_cast(u16[i]) << 16; + std::memcpy(&tmp[i], &v, 4); + } + reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32); + } + + explicit BF16Vec32(const BF16Vec8& v) { + float tmp_small[8]; + RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp_small, v.reg_fp32, 8); + float tmp_large[32]; + for (int i = 0; i < 4; ++i) { + std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float)); + } + reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp_large, 32); + } + + void save(void* ptr) const { + float tmp[32]; + RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < 32; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + + void save(void* ptr, int elem_num) const { + float tmp[32]; + RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < elem_num; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + + void save_strided(void* ptr, ptrdiff_t stride) const { + float tmp[32]; + RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32); + uint8_t* u8 = static_cast(ptr); + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + for (int i = 0; i < 32; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + uint16_t val = static_cast(v >> 16); + *reinterpret_cast(u8 + i * byte_stride) = val; + } + } +}; +#endif + +// ============================================================================ +// FP32 Implementation +// ============================================================================ + +struct FP32Vec4 : public Vec { + constexpr static int VEC_ELEM_NUM = 4; + fixed_fp32x4_t reg; + explicit FP32Vec4(float v) + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(v, VEC_ELEM_NUM)) {}; + explicit FP32Vec4() + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(0.0f, VEC_ELEM_NUM)) {}; + explicit FP32Vec4(const float* ptr) + : reg(RVVI(__riscv_vle32_v_f32, LMUL_128)(ptr, VEC_ELEM_NUM)) {}; + explicit FP32Vec4(fixed_fp32x4_t data) : reg(data) {}; + explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}; + void save(float* ptr) const { + RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, VEC_ELEM_NUM); + } + void save(float* ptr, int elem_num) const { + RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, elem_num); + } +}; + +struct FP32Vec8 : public Vec { + constexpr static int VEC_ELEM_NUM = 8; + fixed_fp32x8_t reg; + + explicit FP32Vec8(float v) + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(v, VEC_ELEM_NUM)) {}; + explicit FP32Vec8() + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, VEC_ELEM_NUM)) {}; + explicit FP32Vec8(const float* ptr) + : reg(RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, VEC_ELEM_NUM)) {}; + explicit FP32Vec8(fixed_fp32x8_t data) : reg(data) {}; + explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {}; + explicit FP32Vec8(const FP16Vec8& v) + : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {}; + explicit FP32Vec8(fixed_fp16x8_t v) + : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {}; + +#ifdef RISCV_BF16_SUPPORT + explicit FP32Vec8(fixed_bf16x8_t v) + : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {}; + explicit FP32Vec8(const BF16Vec8& v) + : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) { + }; +#else + explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {}; +#endif + + float reduce_sum() const { + rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); + scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_256, _f32m1)(reg, scalar, + VEC_ELEM_NUM); + return __riscv_vfmv_f_s_f32m1_f32(scalar); + } + + FP32Vec8 operator*(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec8 operator+(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfadd_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec8 operator-(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfsub_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec8 operator/(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + + FP32Vec8 min(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec8 max(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec8 abs() const { + return FP32Vec8(RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM)); + } + + FP32Vec8 min(const FP32Vec8& b, int elem_num) const { + return FP32Vec8(RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, elem_num)); + } + FP32Vec8 max(const FP32Vec8& b, int elem_num) const { + return FP32Vec8(RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, elem_num)); + } + + FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const { + fixed_fp32x8_t temp = + RVVI(__riscv_vfmax_vv_f32, LMUL_256)(min_v.reg, reg, VEC_ELEM_NUM); + return FP32Vec8( + RVVI(__riscv_vfmin_vv_f32, LMUL_256)(max_v.reg, temp, VEC_ELEM_NUM)); + } + + void save(float* ptr) const { + RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, VEC_ELEM_NUM); + } + void save(float* ptr, int elem_num) const { + RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, elem_num); + } + void save_strided(float* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(float); + RVVI(__riscv_vsse32_v_f32, LMUL_256)(ptr, byte_stride, reg, VEC_ELEM_NUM); + } + + FP32Vec8 exp() const { + const float inv_ln2 = 1.44269504088896341f; + fixed_fp32x8_t x_scaled = + RVVI(__riscv_vfmul_vf_f32, LMUL_256)(reg, inv_ln2, VEC_ELEM_NUM); + fixed_i32x8_t n_int = + RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_256)(x_scaled, VEC_ELEM_NUM); + fixed_fp32x8_t n_float = + RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(n_int, VEC_ELEM_NUM); + + fixed_fp32x8_t r = + RVVI(__riscv_vfsub_vv_f32, LMUL_256)(x_scaled, n_float, VEC_ELEM_NUM); + + fixed_fp32x8_t poly = + RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.001333355810164f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.009618129107628f, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.055504108664821f, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.240226506959101f, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.693147180559945f, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 1.0f, VEC_ELEM_NUM); + + fixed_i32x8_t biased_exp = + RVVI(__riscv_vadd_vx_i32, LMUL_256)(n_int, 127, VEC_ELEM_NUM); + biased_exp = + RVVI(__riscv_vmax_vx_i32, LMUL_256)(biased_exp, 0, VEC_ELEM_NUM); + fixed_i32x8_t exponent_bits = + RVVI(__riscv_vsll_vx_i32, LMUL_256)(biased_exp, 23, VEC_ELEM_NUM); + fixed_fp32x8_t scale = RVVI4(__riscv_vreinterpret_v_i32, LMUL_256, _f32, + LMUL_256)(exponent_bits); + + return FP32Vec8( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, scale, VEC_ELEM_NUM)); + } + + FP32Vec8 tanh() const { + fixed_fp32x8_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_256)( + RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, -9.0f, VEC_ELEM_NUM), 9.0f, + VEC_ELEM_NUM); + fixed_fp32x8_t x2 = + RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x_clamped, 2.0f, VEC_ELEM_NUM); + FP32Vec8 exp_val = FP32Vec8(x2).exp(); + fixed_fp32x8_t num = + RVVI(__riscv_vfsub_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM); + fixed_fp32x8_t den = + RVVI(__riscv_vfadd_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM); + return FP32Vec8( + RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(num, den, VEC_ELEM_NUM)); + } + + FP32Vec8 er() const { + const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f, + a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f; + fixed_fp32x8_t abs_x = + RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM); + + fixed_fp32x8_t t = RVVI(__riscv_vfadd_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vf_f32, LMUL_256)(abs_x, p, VEC_ELEM_NUM), 1.0f, + VEC_ELEM_NUM); + t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_256)(t, 1.0f, VEC_ELEM_NUM); + + fixed_fp32x8_t poly = + RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(a5, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a4, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a3, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a2, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a1, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM); + + fixed_fp32x8_t exp_val = FP32Vec8(RVVI(__riscv_vfneg_v_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)( + abs_x, abs_x, VEC_ELEM_NUM), + VEC_ELEM_NUM)) + .exp() + .reg; + fixed_fp32x8_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, exp_val, VEC_ELEM_NUM), 1.0f, + VEC_ELEM_NUM); + + rvv_mask_f32x8_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_256, BOOL_256)( + reg, 0.0f, VEC_ELEM_NUM); + return FP32Vec8( + RVVI3(__riscv_vfneg_v_f32, LMUL_256, _m)(mask, res, VEC_ELEM_NUM)); + } +}; + +struct FP32Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + fixed_fp32x16_t reg; + + explicit FP32Vec16(float v) + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(v, VEC_ELEM_NUM)) {}; + explicit FP32Vec16() + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.0f, VEC_ELEM_NUM)) {}; + explicit FP32Vec16(const float* ptr) + : reg(RVVI(__riscv_vle32_v_f32, LMUL_512)(ptr, VEC_ELEM_NUM)) {}; + explicit FP32Vec16(fixed_fp32x16_t data) : reg(data) {}; + explicit FP32Vec16(const FP32Vec8& data) + : reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)( + data.reg, data.reg)) {}; + explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {}; + explicit FP32Vec16(const FP16Vec16& v); + +#ifdef RISCV_BF16_SUPPORT + explicit FP32Vec16(fixed_bf16x16_t v) + : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v, VEC_ELEM_NUM)) {}; + explicit FP32Vec16(const BF16Vec16& v) + : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM)) { + }; +#else + explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {}; +#endif + + FP32Vec16 operator+(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec16 operator-(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfsub_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec16 operator*(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec16 operator/(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + + FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfmacc_vv_f32, LMUL_512)(reg, a.reg, b.reg, VEC_ELEM_NUM)); + } + + float reduce_sum() const { + rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); + scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1)(reg, scalar, + VEC_ELEM_NUM); + return __riscv_vfmv_f_s_f32m1_f32(scalar); + } + + float reduce_max() const { + rvv_f32_accum_t scalar = + __riscv_vfmv_s_f_f32m1(std::numeric_limits::lowest(), 1); + scalar = RVVI3(__riscv_vfredmax_vs_f32, LMUL_512, _f32m1)(reg, scalar, + VEC_ELEM_NUM); + return __riscv_vfmv_f_s_f32m1_f32(scalar); + } + + float reduce_min() const { + rvv_f32_accum_t scalar = + __riscv_vfmv_s_f_f32m1(std::numeric_limits::max(), 1); + scalar = RVVI3(__riscv_vfredmin_vs_f32, LMUL_512, _f32m1)(reg, scalar, + VEC_ELEM_NUM); + return __riscv_vfmv_f_s_f32m1_f32(scalar); + } + + template + float reduce_sub_sum(int idx) { + static_assert(VEC_ELEM_NUM % group_size == 0); + const int start = idx * group_size; + auto indices = RVVI(__riscv_vid_v_u32, LMUL_512)(VEC_ELEM_NUM); + rvv_mask_f32x16_t mask = RVVI(__riscv_vmand_mm_, BOOL_512)( + RVVIB(__riscv_vmsgeu_vx_u32, LMUL_512, BOOL_512)(indices, start, + VEC_ELEM_NUM), + RVVIB(__riscv_vmsltu_vx_u32, LMUL_512, BOOL_512)( + indices, start + group_size, VEC_ELEM_NUM), + VEC_ELEM_NUM); + rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); + scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1_m)( + mask, reg, scalar, VEC_ELEM_NUM); + return __riscv_vfmv_f_s_f32m1_f32(scalar); + }; + + FP32Vec16 max(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec16 min(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec16 abs() const { + return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM)); + } + + FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const { + return FP32Vec16(RVVI(__riscv_vfmin_vv_f32, LMUL_512)( + max_v.reg, + RVVI(__riscv_vfmax_vv_f32, LMUL_512)(min_v.reg, reg, VEC_ELEM_NUM), + VEC_ELEM_NUM)); + } + + void save(float* ptr) const { + RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, VEC_ELEM_NUM); + } + void save(float* ptr, int elem_num) const { + RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, elem_num); + } + void save_strided(float* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(float); + RVVI(__riscv_vsse32_v_f32, LMUL_512)(ptr, byte_stride, reg, VEC_ELEM_NUM); + } + + FP32Vec16 exp() const { + const float inv_ln2 = 1.44269504088896341f; + fixed_fp32x16_t x_scaled = + RVVI(__riscv_vfmul_vf_f32, LMUL_512)(reg, inv_ln2, VEC_ELEM_NUM); + fixed_i32x16_t n_int = + RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(x_scaled, VEC_ELEM_NUM); + fixed_fp32x16_t n_float = + RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_512)(n_int, VEC_ELEM_NUM); + fixed_fp32x16_t r = + RVVI(__riscv_vfsub_vv_f32, LMUL_512)(x_scaled, n_float, VEC_ELEM_NUM); + + fixed_fp32x16_t poly = + RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.001333355810164f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), + 0.009618129107628f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), + 0.055504108664821f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), + 0.240226506959101f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), + 0.693147180559945f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), 1.0f, + VEC_ELEM_NUM); + + fixed_i32x16_t biased_exp = RVVI(__riscv_vmax_vx_i32, LMUL_512)( + RVVI(__riscv_vadd_vx_i32, LMUL_512)(n_int, 127, VEC_ELEM_NUM), 0, + VEC_ELEM_NUM); + fixed_fp32x16_t scale = + RVVI4(__riscv_vreinterpret_v_i32, LMUL_512, _f32, LMUL_512)( + RVVI(__riscv_vsll_vx_i32, LMUL_512)(biased_exp, 23, VEC_ELEM_NUM)); + + return FP32Vec16( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, scale, VEC_ELEM_NUM)); + } + + FP32Vec16 tanh() const { + fixed_fp32x16_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_512)( + RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, -9.0f, VEC_ELEM_NUM), 9.0f, + VEC_ELEM_NUM); + FP32Vec16 exp_val = FP32Vec16(RVVI(__riscv_vfmul_vf_f32, LMUL_512)( + x_clamped, 2.0f, VEC_ELEM_NUM)) + .exp(); + return FP32Vec16(RVVI(__riscv_vfdiv_vv_f32, LMUL_512)( + RVVI(__riscv_vfsub_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM), + RVVI(__riscv_vfadd_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM), + VEC_ELEM_NUM)); + } + + FP32Vec16 er() const { + const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f, + a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f; + fixed_fp32x16_t abs_x = + RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM); + fixed_fp32x16_t t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_512)( + RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vf_f32, LMUL_512)(abs_x, p, VEC_ELEM_NUM), 1.0f, + VEC_ELEM_NUM), + 1.0f, VEC_ELEM_NUM); + + fixed_fp32x16_t poly = + RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(a5, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a4, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a3, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a2, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a1, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM); + + fixed_fp32x16_t exp_val = + FP32Vec16(RVVI(__riscv_vfneg_v_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(abs_x, abs_x, + VEC_ELEM_NUM), + VEC_ELEM_NUM)) + .exp() + .reg; + fixed_fp32x16_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, exp_val, VEC_ELEM_NUM), 1.0f, + VEC_ELEM_NUM); + + rvv_mask_f32x16_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_512, BOOL_512)( + reg, 0.0f, VEC_ELEM_NUM); + return FP32Vec16( + RVVI3(__riscv_vfneg_v_f32, LMUL_512, _m)(mask, res, VEC_ELEM_NUM)); + } +}; + +// ============================================================================ +// Type Traits & Global Helpers +// ============================================================================ + +template +struct VecType { + using vec_type = void; + using vec_t = void; +}; + +template +using vec_t = typename VecType::vec_type; + +template <> +struct VecType { + using vec_type = FP32Vec8; + using vec_t = FP32Vec8; +}; +template <> +struct VecType { + using vec_type = FP16Vec8; + using vec_t = FP16Vec8; +}; +template <> +struct VecType { + using vec_type = BF16Vec8; + using vec_t = BF16Vec8; +}; + +template +void storeFP32(float v, T* ptr) { + *ptr = v; +} +template <> +inline void storeFP32(float v, c10::Half* ptr) { + *reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v); +} + +inline FP16Vec16::FP16Vec16(const FP32Vec16& v) { + reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_256)(v.reg, VEC_ELEM_NUM); +} +inline FP16Vec8::FP16Vec8(const FP32Vec8& v) { + reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_128)(v.reg, VEC_ELEM_NUM); +} +inline FP32Vec16::FP32Vec16(const FP16Vec16& v) { + reg = RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM); +} +inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) { + acc = acc.fma(a, b); +} + +#ifdef RISCV_BF16_SUPPORT +template <> +inline void storeFP32(float v, c10::BFloat16* ptr) { + *ptr = static_cast<__bf16>(v); +}; +inline BF16Vec8::BF16Vec8(const FP32Vec8& v) + : reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_128)(v.reg, VEC_ELEM_NUM)) { + }; +inline BF16Vec16::BF16Vec16(const FP32Vec16& v) + : reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_256)(v.reg, VEC_ELEM_NUM)) { + }; +#else +template <> +inline void storeFP32(float v, c10::BFloat16* ptr) { + uint32_t val; + std::memcpy(&val, &v, 4); + *reinterpret_cast(ptr) = static_cast(val >> 16); +} +inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {} +inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {} +#endif + +inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); } + +} // namespace vec_op + +#ifndef CPU_KERNEL_GUARD_IN + #define CPU_KERNEL_GUARD_IN(NAME) +#endif + +#ifndef CPU_KERNEL_GUARD_OUT + #define CPU_KERNEL_GUARD_OUT(NAME) +#endif + +#endif // CPU_TYPES_RISCV_IMPL_HPP From e729cc823d313aa7623ecadefe4305ea241c3dce Mon Sep 17 00:00:00 2001 From: San-Nguyen Date: Mon, 20 Apr 2026 02:25:06 -0500 Subject: [PATCH 131/150] [Fix] Add Spacing when Requesting Output Token > max_model_len (#40324) Signed-off-by: San-Nguyen --- vllm/renderers/params.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index a2c95690c79..10bc01043f7 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -204,7 +204,7 @@ class TokenizeParams: and max_output_tokens > max_total_tokens ): raise VLLMValidationError( - f"{self.max_output_tokens_param}={max_output_tokens}" + f"{self.max_output_tokens_param}={max_output_tokens} " f"cannot be greater than " f"{self.max_total_tokens_param}={max_total_tokens=}. " f"Please request fewer output tokens.", From 77fd2c863147a01c75e1080523425dea68fa75d7 Mon Sep 17 00:00:00 2001 From: wuyingjun <71485611@qq.com> Date: Mon, 20 Apr 2026 15:56:56 +0800 Subject: [PATCH 132/150] [Bugfix] Forward mm_processor_kwargs in offline generate APIs (#40251) Signed-off-by: wuyingjun --- .../llm/test_mm_processor_kwargs.py | 241 ++++++++++++++++++ vllm/entrypoints/llm.py | 42 ++- 2 files changed, 279 insertions(+), 4 deletions(-) create mode 100644 tests/entrypoints/llm/test_mm_processor_kwargs.py diff --git a/tests/entrypoints/llm/test_mm_processor_kwargs.py b/tests/entrypoints/llm/test_mm_processor_kwargs.py new file mode 100644 index 00000000000..19cf91230ca --- /dev/null +++ b/tests/entrypoints/llm/test_mm_processor_kwargs.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from vllm import LLM, SamplingParams + + +def _make_mock_llm() -> LLM: + llm = object.__new__(LLM) + llm.model_config = SimpleNamespace(runner_type="generate") + return llm + + +def test_generate_forwards_mm_processor_kwargs() -> None: + llm = _make_mock_llm() + mm_processor_kwargs = {"num_crops": 4} + sampling_params = SamplingParams(max_tokens=1) + + llm._run_completion = Mock(return_value=["ok"]) + + outputs = llm.generate( + "prompt", + sampling_params=sampling_params, + mm_processor_kwargs=mm_processor_kwargs, + ) + + assert outputs == ["ok"] + assert llm._run_completion.call_args.kwargs["mm_processor_kwargs"] == ( + mm_processor_kwargs + ) + + +def test_enqueue_forwards_mm_processor_kwargs() -> None: + llm = _make_mock_llm() + mm_processor_kwargs = {"do_resize": False} + sampling_params = SamplingParams(max_tokens=1) + + llm._add_completion_requests = Mock(return_value=["req-0"]) + + request_ids = llm.enqueue( + "prompt", + sampling_params=sampling_params, + use_tqdm=False, + mm_processor_kwargs=mm_processor_kwargs, + ) + + assert request_ids == ["req-0"] + assert llm._add_completion_requests.call_args.kwargs["mm_processor_kwargs"] == ( + mm_processor_kwargs + ) + + +def test_chat_forwards_mm_processor_kwargs() -> None: + llm = _make_mock_llm() + mm_processor_kwargs = {"do_pan_and_scan": True} + sampling_params = SamplingParams(max_tokens=1) + messages = [{"role": "user", "content": "hello"}] + + llm._run_chat = Mock(return_value=["ok"]) + + outputs = llm.chat( + messages, + sampling_params=sampling_params, + mm_processor_kwargs=mm_processor_kwargs, + ) + + assert outputs == ["ok"] + assert llm._run_chat.call_args.kwargs["mm_processor_kwargs"] == ( + mm_processor_kwargs + ) + + +def test_run_completion_forwards_mm_processor_kwargs() -> None: + llm = _make_mock_llm() + mm_processor_kwargs = {"min_pixels": 4 * 28 * 28} + sampling_params = SamplingParams(max_tokens=1) + sentinel_output = ["done"] + + llm._add_completion_requests = Mock() + llm._run_engine = Mock(return_value=sentinel_output) + + outputs = llm._run_completion( + prompts=["prompt"], + params=sampling_params, + output_type=object, + use_tqdm=False, + mm_processor_kwargs=mm_processor_kwargs, + ) + + assert outputs == sentinel_output + assert llm._add_completion_requests.call_args.kwargs["mm_processor_kwargs"] == ( + mm_processor_kwargs + ) + + +def test_add_completion_requests_forwards_mm_processor_kwargs() -> None: + llm = _make_mock_llm() + mm_processor_kwargs = {"max_dynamic_patch": 4} + sampling_params = SamplingParams(max_tokens=1) + + llm._params_to_seq = Mock(return_value=[sampling_params]) + llm._lora_request_to_seq = Mock(return_value=[None]) + llm._priority_to_seq = Mock(return_value=[0]) + llm._preprocess_cmpl_one = Mock(return_value={"prompt_token_ids": [1]}) + + captured_prompts = [] + + def fake_render_and_add_requests(*, prompts, **_kwargs): + captured_prompts.extend(prompts) + return ["req-0"] + + llm._render_and_add_requests = Mock(side_effect=fake_render_and_add_requests) + + request_ids = llm._add_completion_requests( + prompts=["prompt"], + params=sampling_params, + use_tqdm=False, + mm_processor_kwargs=mm_processor_kwargs, + ) + + assert request_ids == ["req-0"] + llm._preprocess_cmpl_one.assert_called_once_with( + "prompt", + None, + mm_processor_kwargs=mm_processor_kwargs, + ) + assert captured_prompts == [{"prompt_token_ids": [1]}] + + +def test_preprocess_cmpl_applies_mm_processor_kwargs_to_renderer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + llm = _make_mock_llm() + mm_processor_kwargs = {"num_crops": 8} + prompt = {"prompt": "", "multi_modal_data": {"image": object()}} + + renderer = Mock() + renderer.default_cmpl_tok_params = Mock() + renderer.default_cmpl_tok_params.with_kwargs.return_value = "tok-params" + renderer.render_cmpl.return_value = ["engine-input"] + llm.renderer = renderer + + monkeypatch.setattr( + "vllm.entrypoints.llm.parse_model_prompt", + lambda _model_config, parsed_prompt: parsed_prompt, + ) + + outputs = llm._preprocess_cmpl( + [prompt], + mm_processor_kwargs=mm_processor_kwargs, + ) + + assert outputs == ["engine-input"] + renderer.render_cmpl.assert_called_once_with( + [prompt], + "tok-params", + prompt_extras={"mm_processor_kwargs": mm_processor_kwargs}, + ) + + +def test_preprocess_cmpl_keeps_prompt_mm_processor_kwargs_when_no_override( + monkeypatch: pytest.MonkeyPatch, +) -> None: + llm = _make_mock_llm() + prompt = { + "prompt": "", + "multi_modal_data": {"image": object()}, + "mm_processor_kwargs": {"num_crops": 2}, + } + + renderer = Mock() + renderer.default_cmpl_tok_params = Mock() + renderer.default_cmpl_tok_params.with_kwargs.return_value = "tok-params" + renderer.render_cmpl.return_value = ["engine-input"] + llm.renderer = renderer + + monkeypatch.setattr( + "vllm.entrypoints.llm.parse_model_prompt", + lambda _model_config, parsed_prompt: parsed_prompt, + ) + + outputs = llm._preprocess_cmpl([prompt]) + + assert outputs == ["engine-input"] + renderer.render_cmpl.assert_called_once_with( + [prompt], + "tok-params", + prompt_extras=None, + ) + + +def test_preprocess_chat_applies_mm_processor_kwargs_to_renderer() -> None: + llm = _make_mock_llm() + mm_processor_kwargs = {"num_crops": 8} + messages = [[{"role": "user", "content": "Describe this image."}]] + + renderer = Mock() + renderer.tokenizer = object() + renderer.default_chat_tok_params = Mock() + renderer.default_chat_tok_params.with_kwargs.return_value = "tok-params" + renderer.render_chat.return_value = (messages, ["engine-input"]) + llm.renderer = renderer + + outputs = llm._preprocess_chat( + messages, + mm_processor_kwargs=mm_processor_kwargs, + ) + + assert outputs == ["engine-input"] + call_args = renderer.render_chat.call_args + assert call_args.args[0] == messages + assert call_args.args[1].mm_processor_kwargs == mm_processor_kwargs + assert call_args.args[2] == "tok-params" + assert call_args.kwargs["prompt_extras"] == { + "mm_processor_kwargs": mm_processor_kwargs + } + + +def test_preprocess_chat_omits_mm_processor_kwargs_when_no_override() -> None: + llm = _make_mock_llm() + messages = [[{"role": "user", "content": "Describe this image."}]] + + renderer = Mock() + renderer.tokenizer = object() + renderer.default_chat_tok_params = Mock() + renderer.default_chat_tok_params.with_kwargs.return_value = "tok-params" + renderer.render_chat.return_value = (messages, ["engine-input"]) + llm.renderer = renderer + + outputs = llm._preprocess_chat(messages) + + assert outputs == ["engine-input"] + call_args = renderer.render_chat.call_args + assert call_args.args[0] == messages + assert call_args.args[1].mm_processor_kwargs is None + assert call_args.args[2] == "tok-params" + assert call_args.kwargs["prompt_extras"] is None diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index aaef7528253..98f5e78de41 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -452,6 +452,7 @@ class LLM: lora_request: Sequence[LoRARequest] | LoRARequest | None = None, priority: list[int] | None = None, tokenization_kwargs: dict[str, Any] | None = None, + mm_processor_kwargs: dict[str, Any] | None = None, ) -> list[RequestOutput]: """Generates the completions for the input prompts. @@ -479,6 +480,7 @@ class LLM: of `prompts`, where each priority value corresponds to the prompt at the same index. tokenization_kwargs: Overrides for `tokenizer.encode`. + mm_processor_kwargs: Overrides for `processor.__call__`. Returns: A list of `RequestOutput` objects containing the @@ -503,6 +505,7 @@ class LLM: lora_request=lora_request, tokenization_kwargs=tokenization_kwargs, priority=priority, + mm_processor_kwargs=mm_processor_kwargs, ) def enqueue( @@ -513,6 +516,7 @@ class LLM: priority: list[int] | None = None, use_tqdm: bool | Callable[..., tqdm] = True, tokenization_kwargs: dict[str, Any] | None = None, + mm_processor_kwargs: dict[str, Any] | None = None, ) -> list[str]: """Enqueue prompts for generation without waiting for completion. @@ -527,6 +531,7 @@ class LLM: priority: The priority of the requests, if any. use_tqdm: If True, shows a tqdm progress bar while adding requests. tokenization_kwargs: Overrides for `tokenizer.encode`. + mm_processor_kwargs: Overrides for `processor.__call__`. Returns: A list of request IDs for the enqueued requests. @@ -545,6 +550,7 @@ class LLM: lora_request=lora_request, priority=priority, tokenization_kwargs=tokenization_kwargs, + mm_processor_kwargs=mm_processor_kwargs, ) @overload @@ -843,6 +849,7 @@ class LLM: self, prompts: Sequence[PromptType], tokenization_kwargs: dict[str, Any] | None = None, + mm_processor_kwargs: dict[str, Any] | None = None, ) -> Sequence[EngineInput]: """ Convert prompt inputs from LLM APIs (other than [LLM.chat][]) into @@ -862,15 +869,29 @@ class LLM: tok_params = renderer.default_cmpl_tok_params.with_kwargs( **(tokenization_kwargs or {}) ) + prompt_extras = ( + None + if mm_processor_kwargs is None + else {"mm_processor_kwargs": mm_processor_kwargs} + ) - return renderer.render_cmpl(parsed_prompts, tok_params) + return renderer.render_cmpl( + parsed_prompts, + tok_params, + prompt_extras=prompt_extras, + ) def _preprocess_cmpl_one( self, prompt: PromptType, tokenization_kwargs: dict[str, Any] | None = None, + mm_processor_kwargs: dict[str, Any] | None = None, ) -> EngineInput: - (engine_input,) = self._preprocess_cmpl([prompt], tokenization_kwargs) + (engine_input,) = self._preprocess_cmpl( + [prompt], + tokenization_kwargs, + mm_processor_kwargs=mm_processor_kwargs, + ) return engine_input def _preprocess_chat( @@ -908,16 +929,22 @@ class LLM: tokenize=is_mistral_tokenizer(renderer.tokenizer), ), ), + mm_processor_kwargs=mm_processor_kwargs, ) tok_params = renderer.default_chat_tok_params.with_kwargs( **(tokenization_kwargs or {}) ) + prompt_extras = ( + None + if mm_processor_kwargs is None + else {"mm_processor_kwargs": mm_processor_kwargs} + ) _, engine_inputs = renderer.render_chat( conversations, chat_params, tok_params, - prompt_extras={"mm_processor_kwargs": mm_processor_kwargs}, + prompt_extras=prompt_extras, ) return engine_inputs @@ -1571,6 +1598,7 @@ class LLM: lora_request: Sequence[LoRARequest] | LoRARequest | None = None, priority: list[int] | None = None, tokenization_kwargs: dict[str, Any] | None = None, + mm_processor_kwargs: dict[str, Any] | None = None, ) -> list[str]: seq_prompts = prompt_to_seq(prompts) seq_params = self._params_to_seq(params, len(seq_prompts)) @@ -1579,7 +1607,11 @@ class LLM: return self._render_and_add_requests( prompts=( - self._preprocess_cmpl_one(prompt, tokenization_kwargs) + self._preprocess_cmpl_one( + prompt, + tokenization_kwargs, + mm_processor_kwargs=mm_processor_kwargs, + ) for prompt in maybe_tqdm( seq_prompts, use_tqdm=use_tqdm, @@ -1603,6 +1635,7 @@ class LLM: lora_request: Sequence[LoRARequest] | LoRARequest | None = None, priority: list[int] | None = None, tokenization_kwargs: dict[str, Any] | None = None, + mm_processor_kwargs: dict[str, Any] | None = None, ): self._add_completion_requests( prompts=prompts, @@ -1611,6 +1644,7 @@ class LLM: lora_request=lora_request, priority=priority, tokenization_kwargs=tokenization_kwargs, + mm_processor_kwargs=mm_processor_kwargs, ) return self._run_engine(use_tqdm=use_tqdm, output_type=output_type) From 6d8b80802ba0c6b7c119610af795f674bd78a73f Mon Sep 17 00:00:00 2001 From: milesial Date: Mon, 20 Apr 2026 01:09:44 -0700 Subject: [PATCH 133/150] [Docs] Fix thinking_token_budget docs (#40316) Signed-off-by: milesial --- docs/features/reasoning_outputs.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md index 5c99eb20bc7..c7b2d688f22 100644 --- a/docs/features/reasoning_outputs.md +++ b/docs/features/reasoning_outputs.md @@ -283,9 +283,7 @@ curl http://localhost:8000/v1/chat/completions \ "messages": [ { "role": "user", "content": "9.11 and 9.8, which is greater?" } ], - "extra_body": { - "thinking_token_budget": 10 - } + "thinking_token_budget": 10 }' ``` From a943839e9a7c9ee00e57fcfabbbf0a453393cc97 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 20 Apr 2026 03:09:58 -0500 Subject: [PATCH 134/150] [ROCm][CI] Introducing new MI300 nodes (#39531) Signed-off-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 4955 ++++++++--------- .../{test_mi3xx_moe.py => test_gfx950_moe.py} | 4 +- 2 files changed, 2303 insertions(+), 2656 deletions(-) rename tests/quantization/{test_mi3xx_moe.py => test_gfx950_moe.py} (58%) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index fe617cb10a1..a8c9e409438 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1,8 +1,8 @@ # In this file, you can add more tests to run either by adding a new step or # adding a new command to an existing step. See different options here for examples. -# This script will be feed into Jinja template in `test-template-aws.j2` at -# https://github.com/vllm-project/buildkite-ci/blob/main/scripts/test-template-aws.j2 +# This script will be feed into Jinja template in `test-template-amd.j2` at +# https://github.com/vllm-project/buildkite-ci/blob/main/scripts/test-template-amd.j2 # to generate the final pipeline yaml file. # Documentation @@ -39,8 +39,8 @@ ##################################################################################################################################### # # # IMPORTANT: # -# * Currently AMD CI has MI250 agents, MI325 agents, and MI355 agents. All upcoming feature improvements are tracked in: # -# https://github.com/vllm-project/vllm/issues/34994 # +# * Currently AMD CI has MI250 agents, MI300 agents, MI325 agents, and MI355 agents. All upcoming feature improvements are # +# tracked in: https://github.com/vllm-project/vllm/issues/34994 # # # #-----------------------------------------------------------------------------------------------------------------------------------# # # @@ -104,209 +104,179 @@ ##################################################################################################################################### - - steps: +######################################################################################################################################### +# # +# MI250 (gfx90a) tests # +# # +######################################################################################################################################### -##################################################################################################################################### -# # -# MI250 test definitions ( currently the test set is completely mirrored // TBD which tests are to be routed there ultimately) # -# # -##################################################################################################################################### +#----------------------------------------------------- mi250 · basic_correctness -----------------------------------------------------# -- label: Pytorch Nightly Dependency Override Check # TBD +- label: Distributed Model Tests (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_2 + num_gpus: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/model_executor/model_loader/sharded_state_loader.py + - vllm/model_executor/models/ + - vllm/model_executor/layers/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - tests/basic_correctness/ + - tests/model_executor/model_loader/test_sharded_state_loader.py + - tests/models/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' + - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/language -v -s -m 'distributed(num_gpus=2)' + - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py + - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' + +#-------------------------------------------------------- mi250 · benchmarks ---------------------------------------------------------# + +- label: Benchmarks # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + working_dir: "/vllm-workspace/.buildkite" + source_file_dependencies: + - benchmarks/ + - vllm/platforms/rocm.py + commands: + - bash scripts/run-benchmarks.sh + +#---------------------------------------------------------- mi250 · compile ----------------------------------------------------------# + +- label: PyTorch Compilation Unit Tests # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + torch_nightly: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/compilation/ + - vllm/model_executor/layers/ + - vllm/v1/worker/ + - vllm/v1/attention/ + - vllm/v1/cudagraph_dispatcher.py + - vllm/config/compilation.py + - csrc/ + - tests/compile + - vllm/platforms/rocm.py + commands: + - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" + +- label: PyTorch Fullgraph # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true - soft_fail: true + torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - requirements/test/nightly-torch.txt + - vllm/compilation/ + - vllm/model_executor/ + - vllm/v1/attention/ + - vllm/config/compilation.py + - csrc/ + - tests/compile - vllm/platforms/rocm.py commands: - - bash standalone_tests/pytorch_nightly_dependency.sh + - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile' - -- label: Async Engine, Inputs, Utils, Worker # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/detokenizer - - tests/multimodal - - tests/utils_ - commands: - - pytest -v -s detokenizer - - pytest -v -s -m 'not cpu_test' multimodal - - pytest -v -s utils_ - - -- label: Async Engine, Inputs, Utils, Worker, Config (CPU) # TBD +- label: PyTorch Fullgraph Smoke Test # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true - no_gpu: true + torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/test_inputs.py - - tests/test_outputs.py - - tests/test_pooling_params.py - - tests/test_ray_env.py - - tests/multimodal - - tests/renderers - - tests/standalone_tests/lazy_imports.py - - tests/tokenizers_ - - tests/tool_parsers - - tests/transformers_utils - - tests/config + - vllm/compilation/ + - vllm/model_executor/ + - vllm/v1/attention/ + - vllm/config/compilation.py + - csrc/ + - tests/compile + - vllm/platforms/rocm.py commands: - - python3 standalone_tests/lazy_imports.py - - pytest -v -s test_inputs.py - - pytest -v -s test_outputs.py - - pytest -v -s test_pooling_params.py - - pytest -v -s test_ray_env.py - - pytest -v -s -m 'cpu_test' multimodal - - pytest -v -s renderers - - pytest -v -s tokenizers_ - - pytest -v -s tool_parsers - - pytest -v -s transformers_utils - - pytest -v -s config + - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\\\;" - -- label: Python-only Installation # TBD +- label: Distributed Compile + RPC Tests (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 + agent_pool: mi250_2 + num_gpus: 2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - tests/standalone_tests/python_only_compile.sh - - setup.py + - vllm/compilation/ + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/compile/fullgraph/test_basic_correctness.py + - tests/compile/test_wrapper.py + - tests/entrypoints/llm/test_collective_rpc.py - vllm/platforms/rocm.py commands: - - bash standalone_tests/python_only_compile.sh + - export TORCH_NCCL_BLOCKING_WAIT=1 + - pytest -v -s entrypoints/llm/test_collective_rpc.py + - pytest -v -s ./compile/fullgraph/test_basic_correctness.py + - pytest -v -s ./compile/test_wrapper.py +#-------------------------------------------------------- mi250 · distributed --------------------------------------------------------# -- label: Basic Correctness # TBD +- label: Distributed Comm Ops # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - fast_check: true - torch_nightly: true + agent_pool: mi250_2 + num_gpus: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/basic_correctness/test_basic_correctness - - tests/basic_correctness/test_cpu_offload - - tests/basic_correctness/test_cumem.py - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s basic_correctness/test_cumem.py - - pytest -v -s basic_correctness/test_basic_correctness.py - - pytest -v -s basic_correctness/test_cpu_offload.py - - -- label: Entrypoints Unit Tests # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - fast_check: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/entrypoints - - tests/entrypoints/ + - vllm/distributed + - tests/distributed - vllm/platforms/rocm.py commands: - - pytest -v -s entrypoints/openai/tool_parsers - - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling + - pytest -v -s distributed/test_comm_ops.py + - pytest -v -s distributed/test_shm_broadcast.py + - pytest -v -s distributed/test_shm_buffer.py + - pytest -v -s distributed/test_shm_storage.py - -- label: Entrypoints Integration (LLM) # TBD +- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - fast_check: true - torch_nightly: true + agent_pool: mi250_2 + num_gpus: 2 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/entrypoints/llm - - tests/entrypoints/offline_mode - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py - - pytest -v -s entrypoints/llm/test_generate.py - - pytest -v -s entrypoints/offline_mode - - -- label: Entrypoints Integration (API Server 2) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/rpc - - tests/entrypoints/serve/instrumentator - - tests/tool_use - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve/instrumentator - - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc - - pytest -v -s tool_use - - -- label: Entrypoints Integration (Responses API) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/openai/responses - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/responses - - -- label: EPLB Algorithm # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/eplb - - tests/distributed/test_eplb_algo.py + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/distributed/ + - tests/v1/shutdown + - tests/v1/worker/test_worker_memory_snapshot.py - vllm/platforms/rocm.py commands: - - pytest -v -s distributed/test_eplb_algo.py - - -- label: EPLB Execution # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/eplb - - tests/distributed/test_eplb_execute.py - - tests/distributed/test_eplb_spec_decode.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s distributed/test_eplb_execute.py - - pytest -v -s distributed/test_eplb_spec_decode.py - + - export TORCH_NCCL_BLOCKING_WAIT=1 + - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' + - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown + - pytest -v -s v1/worker/test_worker_memory_snapshot.py - label: Elastic EP Scaling Test # TBD timeout_in_minutes: 180 @@ -324,37 +294,43 @@ steps: commands: - pytest -v -s distributed/test_elastic_ep.py - -- label: Metrics, Tracing (2 GPUs) # TBD +- label: EPLB Execution # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - num_gpus: 2 + agent_pool: mi250_4 + num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/v1/tracing + - vllm/distributed/eplb + - tests/distributed/test_eplb_execute.py + - tests/distributed/test_eplb_spec_decode.py + - vllm/platforms/rocm.py commands: - - "pip install \ - 'opentelemetry-sdk>=1.26.0' \ - 'opentelemetry-api>=1.26.0' \ - 'opentelemetry-exporter-otlp>=1.26.0' \ - 'opentelemetry-semantic-conventions-ai>=0.4.1'" - - pytest -v -s v1/tracing + - pytest -v -s distributed/test_eplb_execute.py + - pytest -v -s distributed/test_eplb_spec_decode.py - -- label: Regression # TBD +- label: Pipeline + Context Parallelism (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 + agent_pool: mi250_4 + num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/test_regression + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/model_executor/models/ + - vllm/model_executor/layers/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - tests/distributed/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - pip install modelscope - - pytest -v -s test_regression.py + - pytest -v -s distributed/test_pp_cudagraph.py + - pytest -v -s distributed/test_pipeline_parallel.py +#---------------------------------------------------------- mi250 · engine -----------------------------------------------------------# - label: Engine # TBD timeout_in_minutes: 180 @@ -372,200 +348,24 @@ steps: commands: - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py +#----------------------------------------------------------- mi250 · evals -----------------------------------------------------------# -- label: Engine (1 GPU) # TBD +- label: Multi-Modal Accuracy Eval (Small Models) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true - working_dir: "/vllm-workspace/tests" + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" source_file_dependencies: - - vllm/v1/ - - tests/v1/engine/ + - vllm/multimodal/ + - vllm/inputs/ + - vllm/v1/core/ - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/engine/test_preprocess_error_handling.py - - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py - - -- label: e2e Scheduling (1 GPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/ - - tests/v1/e2e/general/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/general/test_async_scheduling.py - - -- label: e2e Core (1 GPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/ - - tests/v1/e2e/general/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py - - -- label: Spec Decode Speculators + MTP # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - vllm/transformers_utils/configs/speculators/ - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" - - -- label: Spec Decode Ngram + Suffix # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix" - - -- label: Spec Decode Draft Model # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" - - -- label: V1 e2e (2 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/e2e - commands: - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" - - -- label: V1 Sample + Logits # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/sample - - tests/v1/logits_processors - - tests/v1/test_oracle.py - - tests/v1/test_request.py - - tests/v1/test_outputs.py - commands: - - pytest -v -s v1/sample - - pytest -v -s v1/logits_processors - - pytest -v -s v1/test_oracle.py - - pytest -v -s v1/test_request.py - - pytest -v -s v1/test_outputs.py - - -- label: V1 Core + KV + Metrics # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/core - - tests/v1/executor - - tests/v1/kv_offload - - tests/v1/worker - - tests/v1/kv_connector/unit - - tests/v1/metrics - - tests/entrypoints/openai/correctness/test_lmeval.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - pytest -v -s -m 'not cpu_test' v1/core - - pytest -v -s v1/executor - - pytest -v -s v1/kv_offload - - pytest -v -s v1/worker - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'not cpu_test' v1/metrics - - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api - - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - - -- label: V1 attention (H100-MI250) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/config/attention.py - - vllm/model_executor/layers/attention - - vllm/v1/attention - - tests/v1/attention - - vllm/_aiter_ops.py - - vllm/envs.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/attention - - -- label: V1 others (CPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - no_gpu: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1 - commands: - - pytest -v -s -m 'cpu_test' v1/core - - pytest -v -s v1/structured_output - - pytest -v -s v1/test_serial_utils.py - - pytest -v -s -m 'cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'cpu_test' v1/metrics + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1 +#--------------------------------------------------------- mi250 · examples ----------------------------------------------------------# - label: Examples # TBD timeout_in_minutes: 180 @@ -601,129 +401,7 @@ steps: - python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048 - python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536 - -- label: Platform Tests (CUDA) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/cuda - commands: - - pytest -v -s cuda/test_cuda_context.py - - pytest -v -s cuda/test_platform_no_cuda_init.py - - -- label: Samplers Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/model_executor/layers - - vllm/sampling_metadata.py - - vllm/v1/sample/ - - vllm/beam_search.py - - tests/samplers - - tests/conftest.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s samplers - - -- label: LoRA %N # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - parallelism: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/lora - - tests/lora - - vllm/platforms/rocm.py - commands: - - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py - - -- label: PyTorch Compilation Unit Tests # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - torch_nightly: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers/ - - vllm/v1/worker/ - - vllm/v1/attention/ - - vllm/v1/cudagraph_dispatcher.py - - vllm/config/compilation.py - - csrc/ - - tests/compile - - vllm/platforms/rocm.py - commands: - - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" - - -- label: PyTorch Fullgraph Smoke Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/ - - vllm/v1/attention/ - - vllm/config/compilation.py - - csrc/ - - tests/compile - - vllm/platforms/rocm.py - commands: - - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\\\;" - - -- label: PyTorch Fullgraph # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/ - - vllm/v1/attention/ - - vllm/config/compilation.py - - csrc/ - - tests/compile - - vllm/platforms/rocm.py - commands: - - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile' - - -- label: Cudagraph # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - tests/v1/cudagraph - - vllm/v1/cudagraph_dispatcher.py - - vllm/config/compilation.py - - vllm/compilation - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py - - pytest -v -s v1/cudagraph/test_cudagraph_mode.py - +#---------------------------------------------------------- mi250 · kernels ----------------------------------------------------------# - label: Kernels Core Operation Test # TBD timeout_in_minutes: 180 @@ -740,23 +418,7 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s kernels/core kernels/test_top_k_per_row.py - - -- label: Kernels Mamba Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/mamba/ - - tests/kernels/mamba - - vllm/model_executor/layers/mamba/ops - - vllm/platforms/rocm.py - commands: - - pytest -v -s kernels/mamba - + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py - label: Kernels Helion Test # TBD timeout_in_minutes: 180 @@ -772,6 +434,37 @@ steps: - pip install helion==1.0.0 - pytest -v -s kernels/helion/ +- label: Kernels Mamba Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/mamba/ + - tests/kernels/mamba + - vllm/model_executor/layers/mamba/ops + - vllm/platforms/rocm.py + commands: + - pytest -v -s kernels/mamba + +#----------------------------------------------------------- mi250 · lora ------------------------------------------------------------# + +- label: LoRA %N # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + parallelism: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/lora + - tests/lora + - vllm/platforms/rocm.py + commands: + - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py + +#------------------------------------------------------ mi250 · model_executor -------------------------------------------------------# - label: Model Executor # TBD timeout_in_minutes: 180 @@ -794,64 +487,22 @@ steps: - pytest -v -s model_executor -m '(not slow_test)' - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py +#---------------------------------------------------------- mi250 · models -----------------------------------------------------------# -- label: Benchmarks # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/.buildkite" - source_file_dependencies: - - benchmarks/ - - vllm/platforms/rocm.py - commands: - - bash scripts/run-benchmarks.sh - - -- label: Benchmarks CLI Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/benchmarks/ - commands: - - pytest -v -s benchmarks/ - - -- label: OpenAI API correctness # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - vllm/entrypoints/openai/ - - vllm/model_executor/models/whisper.py - - vllm/model_executor/layers/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/model_executor/model_loader/ - commands: - - bash ../tools/install_torchcodec_rocm.sh || exit 1 - - pytest -s entrypoints/openai/correctness/ - - -- label: Basic Models Tests (Initialization) # TBD +- label: Basic Models Test (Other CPU) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 + no_gpu: true + optional: true torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/test_initialization.py - - tests/models/registry.py + - tests/models/test_utils.py + - tests/models/test_vision.py commands: - - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset - + - pytest -v -s models/test_utils.py models/test_vision.py - label: Basic Models Tests (Extra Initialization) %N # TBD timeout_in_minutes: 180 @@ -870,6 +521,18 @@ steps: commands: - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB +- label: Basic Models Tests (Initialization) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/test_initialization.py + - tests/models/registry.py + commands: + - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset - label: Basic Models Tests (Other) # TBD timeout_in_minutes: 180 @@ -885,22 +548,27 @@ steps: commands: - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py - -- label: Basic Models Test (Other CPU) # TBD +- label: Language Models Test (MTEB) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - no_gpu: true - optional: true - torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/test_utils.py - - tests/models/test_vision.py + - tests/models/language/pooling_mteb_test commands: - - pytest -v -s models/test_utils.py models/test_vision.py + - pytest -v -s models/language/pooling_mteb_test +- label: Language Models Test (PPL) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/language/generation_ppl_test + commands: + - pytest -v -s models/language/generation_ppl_test - label: Language Models Tests (Extra Standard) %N # TBD timeout_in_minutes: 180 @@ -925,104 +593,28 @@ steps: - export TORCH_NCCL_BLOCKING_WAIT=1 - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - -- label: Language Models Test (PPL) # TBD +- label: Multi-Modal Models (Extended Generation 2) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/language/generation_ppl_test - commands: - - pytest -v -s models/language/generation_ppl_test - - -- label: Language Models Test (Extended Pooling) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/language/pooling - commands: - - pytest -v -s models/language/pooling -m 'not core_model' - - -- label: Language Models Test (MTEB) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/language/pooling_mteb_test - commands: - - pytest -v -s models/language/pooling_mteb_test - - -- label: Multi-Modal Processor (CPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - no_gpu: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - - tests/models/registry.py + - tests/models/multimodal/generation commands: - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py + - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' - -- label: Multi-Modal Accuracy Eval (Small Models) # TBD +- label: Multi-Modal Models (Extended Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" - source_file_dependencies: - - vllm/multimodal/ - - vllm/inputs/ - - vllm/v1/core/ - - vllm/platforms/rocm.py - - vllm/model_executor/model_loader/ - commands: - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1 - - -- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/multimodal + - tests/models/multimodal/pooling commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" - - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model - - -- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" - - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model - + - pytest -v -s models/multimodal/pooling -m 'not core_model' - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" # TBD timeout_in_minutes: 180 @@ -1038,93 +630,221 @@ steps: - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model +#---------------------------------------------------------- mi250 · plugins ----------------------------------------------------------# -- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD +- label: Plugin Tests (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - - tests/models/multimodal/test_mapping.py - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing - - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model - - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model - - -- label: Multi-Modal Models (Extended Generation 1) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - - tests/models/multimodal/test_mapping.py - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py - - pytest -v -s models/multimodal/test_mapping.py - - -- label: Multi-Modal Models (Extended Generation 2) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' - - -- label: Multi-Modal Models (Extended Generation 3) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' - - -- label: Multi-Modal Models (Extended Pooling) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/pooling - commands: - - pytest -v -s models/multimodal/pooling -m 'not core_model' - - -- label: Distributed Comm Ops # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 + agent_pool: mi250_2 num_gpus: 2 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed - - tests/distributed + - vllm/plugins/ + - tests/plugins/ - vllm/platforms/rocm.py commands: - - pytest -v -s distributed/test_comm_ops.py - - pytest -v -s distributed/test_shm_broadcast.py - - pytest -v -s distributed/test_shm_buffer.py - - pytest -v -s distributed/test_shm_storage.py + # BEGIN: platform plugin and general plugin tests, all the code in-between runs on dummy platform + - pip install -e ./plugins/vllm_add_dummy_platform + - pytest -v -s plugins_tests/test_platform_plugins.py + - pip uninstall vllm_add_dummy_platform -y + # END: platform plugin tests + # BEGIN: `io_processor` plugins test, all the code in between uses the `prithvi_io_processor` plugin + - pip install -e ./plugins/prithvi_io_processor_plugin + - pytest -v -s plugins_tests/test_io_processor_plugins.py + - pytest -v -s plugins_tests/test_terratorch_io_processor_plugins.py + - pip uninstall prithvi_io_processor_plugin -y + # END: `io_processor` plugins test + # BEGIN: `bge_m3_sparse io_processor` test + - pip install -e ./plugins/bge_m3_sparse_plugin + - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py + - pip uninstall bge_m3_sparse_plugin -y + # END: `bge_m3_sparse io_processor` test + # BEGIN: `stat_logger` plugins test + - pip install -e ./plugins/vllm_add_dummy_stat_logger + - pytest -v -s plugins_tests/test_stats_logger_plugins.py + - pip uninstall dummy_stat_logger -y + # END: `stat_logger` plugins test + # BEGIN: other tests + - pytest -v -s plugins_tests/test_scheduler_plugins.py + - pip install -e ./plugins/vllm_add_dummy_model + - pytest -v -s distributed/test_distributed_oot.py + - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py + - pytest -v -s models/test_oot_registration.py + - pytest -v -s plugins/lora_resolvers +#------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------# + +- label: Batch Invariance (H100-MI250) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + +- label: Cudagraph # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - tests/v1/cudagraph + - vllm/v1/cudagraph_dispatcher.py + - vllm/config/compilation.py + - vllm/compilation + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py + - pytest -v -s v1/cudagraph/test_cudagraph_mode.py + +- label: e2e Core (1 GPU) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/ + - tests/v1/e2e/general/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py + +- label: e2e Scheduling (1 GPU) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/ + - tests/v1/e2e/general/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/general/test_async_scheduling.py + +- label: Engine (1 GPU) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/ + - tests/v1/engine/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/engine/test_preprocess_error_handling.py + - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py + +- label: Spec Decode Draft Model # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" + +- label: Spec Decode Speculators + MTP # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - vllm/transformers_utils/configs/speculators/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" + +- label: V1 attention (H100-MI250) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/config/attention.py + - vllm/model_executor/layers/attention + - vllm/v1/attention + - tests/v1/attention + - vllm/_aiter_ops.py + - vllm/envs.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/attention + +- label: V1 Core + KV + Metrics # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - pytest -v -s -m 'not cpu_test' v1/core + - pytest -v -s v1/executor + - pytest -v -s v1/kv_offload + - pytest -v -s v1/worker + - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'not cpu_test' v1/metrics + - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api + - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine + +- label: V1 Sample + Logits # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/sample + - tests/v1/logits_processors + - tests/v1/test_oracle.py + - tests/v1/test_request.py + - tests/v1/test_outputs.py + commands: + - pytest -v -s v1/sample + - pytest -v -s v1/logits_processors + - pytest -v -s v1/test_oracle.py + - pytest -v -s v1/test_request.py + - pytest -v -s v1/test_outputs.py - label: Distributed DP Tests (2 GPUs) # TBD timeout_in_minutes: 180 @@ -1150,192 +870,6 @@ steps: - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py - -- label: Distributed Compile + RPC Tests (2 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/compilation/ - - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/worker/worker_base.py - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/compile/fullgraph/test_basic_correctness.py - - tests/compile/test_wrapper.py - - tests/entrypoints/llm/test_collective_rpc.py - - vllm/platforms/rocm.py - commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - pytest -v -s entrypoints/llm/test_collective_rpc.py - - pytest -v -s ./compile/fullgraph/test_basic_correctness.py - - pytest -v -s ./compile/test_wrapper.py - - -- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/worker/worker_base.py - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/distributed/ - - tests/v1/shutdown - - tests/v1/worker/test_worker_memory_snapshot.py - - vllm/platforms/rocm.py - commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - - pytest -v -s v1/worker/test_worker_memory_snapshot.py - - -- label: Distributed Model Tests (2 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_2 - num_gpus: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/model_executor/model_loader/sharded_state_loader.py - - vllm/model_executor/models/ - - vllm/model_executor/layers/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - tests/basic_correctness/ - - tests/model_executor/model_loader/test_sharded_state_loader.py - - tests/models/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' - - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' - - pytest models/language -v -s -m 'distributed(num_gpus=2)' - - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py - - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' - - -- label: Plugin Tests (2 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/plugins/ - - tests/plugins/ - - vllm/platforms/rocm.py - commands: - # BEGIN: platform plugin and general plugin tests, all the code in-between runs on dummy platform - - pip install -e ./plugins/vllm_add_dummy_platform - - pytest -v -s plugins_tests/test_platform_plugins.py - - pip uninstall vllm_add_dummy_platform -y - # END: platform plugin tests - # BEGIN: `io_processor` plugins test, all the code in between uses the `prithvi_io_processor` plugin - - pip install -e ./plugins/prithvi_io_processor_plugin - - pytest -v -s plugins_tests/test_io_processor_plugins.py - - pip uninstall prithvi_io_processor_plugin -y - # END: `io_processor` plugins test - # BEGIN: `bge_m3_sparse io_processor` test - - pip install -e ./plugins/bge_m3_sparse_plugin - - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py - - pip uninstall bge_m3_sparse_plugin -y - # END: `bge_m3_sparse io_processor` test - # BEGIN: `stat_logger` plugins test - - pip install -e ./plugins/vllm_add_dummy_stat_logger - - pytest -v -s plugins_tests/test_stats_logger_plugins.py - - pip uninstall dummy_stat_logger -y - # END: `stat_logger` plugins test - # BEGIN: other tests - - pytest -v -s plugins_tests/test_scheduler_plugins.py - - pip install -e ./plugins/vllm_add_dummy_model - - pytest -v -s distributed/test_distributed_oot.py - - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py - - pytest -v -s models/test_oot_registration.py - - pytest -v -s plugins/lora_resolvers - # END: other tests - - -- label: Pipeline + Context Parallelism (4 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/model_executor/models/ - - vllm/model_executor/layers/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - tests/distributed/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s distributed/test_pp_cudagraph.py - - pytest -v -s distributed/test_pipeline_parallel.py - - -- label: Ray Dependency Compatibility Check # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/" - source_file_dependencies: - - requirements/ - - setup.py - - vllm/platforms/rocm.py - commands: - - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh - - -- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - - -- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - - - label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] @@ -1351,8 +885,19 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh +- label: V1 e2e (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/e2e + commands: + - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" -- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD +- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_4 @@ -1364,28 +909,256 @@ steps: - vllm/platforms/rocm.py commands: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - CROSS_LAYERS_BLOCKS=True ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +#------------------------------------------------------------- mi250 · misc ------------------------------------------------------------# -- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD +- label: Async Engine, Inputs, Utils, Worker, Config (CPU) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_4 - num_gpus: 4 + agent_pool: mi250_1 + optional: true + no_gpu: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ + - vllm/ + - tests/test_inputs.py + - tests/test_outputs.py + - tests/test_pooling_params.py + - tests/test_ray_env.py + - tests/multimodal + - tests/renderers + - tests/standalone_tests/lazy_imports.py + - tests/tokenizers_ + - tests/tool_parsers + - tests/transformers_utils + - tests/config + commands: + - python3 standalone_tests/lazy_imports.py + - pytest -v -s test_inputs.py + - pytest -v -s test_outputs.py + - pytest -v -s test_pooling_params.py + - pytest -v -s test_ray_env.py + - pytest -v -s -m 'cpu_test' multimodal + - pytest -v -s renderers + - pytest -v -s tokenizers_ + - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py --ignore=reasoning/test_gemma4_reasoning_parser.py + - pytest -v -s tool_parsers + - pytest -v -s transformers_utils + - pytest -v -s config + +######################################################################################################################################### +# # +# MI300 (gfx942) tests # +# # +######################################################################################################################################### + +#----------------------------------------------------- mi300 · basic_correctness -----------------------------------------------------# + +- label: Basic Correctness # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/basic_correctness/test_basic_correctness + - tests/basic_correctness/test_cpu_offload + - tests/basic_correctness/test_cumem.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s basic_correctness/test_cumem.py + - pytest -v -s basic_correctness/test_basic_correctness.py + - pytest -v -s basic_correctness/test_cpu_offload.py + +- label: Distributed Model Tests (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/model_executor/model_loader/sharded_state_loader.py + - vllm/model_executor/models/ + - vllm/model_executor/layers/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - tests/basic_correctness/ + - tests/model_executor/model_loader/test_sharded_state_loader.py + - tests/models/ + commands: + - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' + - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/language -v -s -m 'distributed(num_gpus=2)' + - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py + - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' + - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' + +#-------------------------------------------------------- mi300 · benchmarks ---------------------------------------------------------# + +- label: Benchmarks # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/.buildkite" + source_file_dependencies: + - benchmarks/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - HYBRID_SSM=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + - bash scripts/run-benchmarks.sh - -- label: Distributed Tests (2 GPUs)(H100-MI250) # TBD +- label: Benchmarks CLI Test # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi325_2 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/benchmarks/ + commands: + - pytest -v -s benchmarks/ + +#---------------------------------------------------------- mi300 · compile ----------------------------------------------------------# + +- label: Fusion E2E Config Sweep (H100-MI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + num_gpus: 1 + working_dir: "/vllm-workspace/" + source_file_dependencies: + - csrc/quantization/ + - vllm/compilation/ + - vllm/model_executor/layers/layernorm.py + - vllm/model_executor/layers/activation.py + - vllm/model_executor/layers/attention/attention.py + - vllm/model_executor/layers/quantization/input_quant_fp8.py + - tests/compile/fusions_e2e/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - rocm-smi + - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3" + +- label: Fusion E2E Quick (H100-MI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + num_gpus: 1 + working_dir: "/vllm-workspace/" + source_file_dependencies: + - csrc/quantization/ + - vllm/model_executor/ + - vllm/v1/attention/ + - vllm/compilation/ + - tests/compile/fusions_e2e/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - rocm-smi + # Run all models and attn backends but only Inductor partition and native custom ops + - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and not +rms_norm and not +quant_fp8'" + # Different from CUDA, Qwen requires +rms_norm and +quant_fp8 as rms+quant fusion is only supported on AITER + - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and +rms_norm and +quant_fp8 and qwen3'" + +- label: PyTorch Compilation Passes Unit Tests # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/compile/passes + commands: + - pytest -s -v compile/passes --ignore compile/passes/distributed + +- label: Pytorch Nightly Dependency Override Check # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + soft_fail: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - requirements/test/nightly-torch.txt + - vllm/platforms/rocm.py + commands: + - bash standalone_tests/pytorch_nightly_dependency.sh + +- label: Distributed Compile Unit Tests (2xH100-2xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + working_dir: "/vllm-workspace/" + source_file_dependencies: + - vllm/compilation/ + - vllm/model_executor/layers + - tests/compile/passes/distributed/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py + - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py + +#----------------------------------------------------------- mi300 · cuda ------------------------------------------------------------# + +- label: Platform Tests (CUDA) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/cuda + commands: + - pytest -v -s cuda/test_cuda_context.py + - pytest -v -s cuda/test_platform_no_cuda_init.py + +#-------------------------------------------------------- mi300 · detokenizer --------------------------------------------------------# + +- label: Async Engine, Inputs, Utils, Worker # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/detokenizer + - tests/multimodal + - tests/utils_ + commands: + - pytest -v -s detokenizer + - pytest -v -s -m 'not cpu_test' multimodal + - pytest -v -s utils_ + +#-------------------------------------------------------- mi300 · distributed --------------------------------------------------------# + +- label: EPLB Algorithm # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/eplb + - tests/distributed/test_eplb_algo.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s distributed/test_eplb_algo.py + - pytest -v -s distributed/test_eplb_utils.py + +- label: Distributed Tests (2xH100-2xMI250) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 num_gpus: 2 working_dir: "/vllm-workspace/" source_file_dependencies: @@ -1403,122 +1176,26 @@ steps: - pytest -v -s tests/distributed/test_context_parallel.py - VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization - -##################################################################################################################################### -# # -# gfx942 # -# # -##################################################################################################################################### - - -- label: Entrypoints Integration (LLM) # 13.1m - timeout_in_minutes: 22 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Distributed Tests (4xA100-4xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 optional: true - fast_check: true - torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/entrypoints/llm - - tests/entrypoints/offline_mode commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py - - pytest -v -s entrypoints/llm/test_generate.py - - pytest -v -s entrypoints/offline_mode - - -- label: Entrypoints Integration (API Server openai - Part 1) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/openai - - tests/entrypoints/test_chat_utils - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py - - -- label: Entrypoints Integration (API Server openai - Part 2) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/openai - - tests/entrypoints/test_chat_utils - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - - pytest -v -s entrypoints/openai/speech_to_text/ - - pytest -v -s entrypoints/test_chat_utils.py - - -- label: Entrypoints Integration (API Server openai - Part 3) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/openai - - tests/entrypoints/test_chat_utils - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses - - -- label: Entrypoints Integration (API Server 2) #26.9m - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/rpc - - tests/entrypoints/serve/instrumentator - - tests/tool_use - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve/instrumentator - - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc - - pytest -v -s tool_use - - -- label: Entrypoints Integration (Pooling) # 22.8m - timeout_in_minutes: 48 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/pooling - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/pooling - + - export TORCH_NCCL_BLOCKING_WAIT=1 + - pytest -v -s distributed/test_custom_all_reduce.py + - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py + - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - pytest -v -s -x lora/test_mixtral.py - label: Distributed Torchrun + Examples (4 GPUs) # TBD - timeout_in_minutes: 80 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -1541,57 +1218,50 @@ steps: - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_nccl.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_ipc.py - -- label: Distributed DP Tests (4 GPUs) # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 +- label: Elastic EP Scaling Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 num_gpus: 4 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/distributed/ - - tests/v1/distributed - - tests/v1/engine/test_engine_core_client.py - - tests/distributed/test_utils + - vllm/engine/ + - vllm/executor/ + - vllm/compilation/ + - tests/distributed/ - vllm/platforms/rocm.py commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py - - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py - - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py - - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py - - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp - - pytest -v -s distributed/test_utils.py + - pytest -v -s distributed/test_elastic_ep.py - -- label: Distributed Compile + Comm (4 GPUs) # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 +- label: RayExecutorV2 (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/ - - tests/distributed/test_pynccl - - tests/distributed/test_events - - tests/compile/fullgraph/test_basic_correctness.py - - tests/distributed/test_symm_mem_allreduce.py - - tests/distributed/test_multiproc_executor.py + - vllm/v1/executor/ray_executor_v2.py + - vllm/v1/executor/abstract.py + - vllm/v1/executor/multiproc_executor.py + - tests/distributed/test_ray_v2_executor.py + - tests/distributed/test_ray_v2_executor_e2e.py + - tests/distributed/test_pipeline_parallel.py + - tests/basic_correctness/test_basic_correctness.py - vllm/platforms/rocm.py commands: + - export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 - export TORCH_NCCL_BLOCKING_WAIT=1 - - pytest -v -s compile/fullgraph/test_basic_correctness.py - - pytest -v -s distributed/test_pynccl.py - - pytest -v -s distributed/test_events.py - - pytest -v -s distributed/test_symm_mem_allreduce.py - - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node + - pytest -v -s distributed/test_ray_v2_executor.py + - pytest -v -s distributed/test_ray_v2_executor_e2e.py + - pytest -v -s distributed/test_pipeline_parallel.py -k "ray" + - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray" - -- label: Distributed Tests (8 GPUs)(H100-MI325) # 6.4m - timeout_in_minutes: 10 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_8 +- label: Distributed Tests (8xH100-8xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_8 num_gpus: 8 optional: true working_dir: "/vllm-workspace/tests" @@ -1607,331 +1277,359 @@ steps: - export TORCH_NCCL_BLOCKING_WAIT=1 - torchrun --nproc-per-node=8 ../examples/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep +#-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------# -- label: Elastic EP Scaling Test # TBD +- label: Entrypoints Integration (API Server 2) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/entrypoints/rpc + - tests/entrypoints/serve/instrumentator + - tests/tool_use + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/serve/instrumentator + - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc + - pytest -v -s tool_use + +- label: Entrypoints Integration (API Server openai - Part 1) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/entrypoints/openai + - tests/entrypoints/test_chat_utils + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py + +- label: Entrypoints Integration (API Server openai - Part 2) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/entrypoints/openai + - tests/entrypoints/test_chat_utils + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py + - pytest -v -s entrypoints/openai/speech_to_text/ + - pytest -v -s entrypoints/test_chat_utils.py + +- label: Entrypoints Integration (API Server openai - Part 3) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/entrypoints/openai + - tests/entrypoints/test_chat_utils + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py + +- label: Entrypoints Integration (LLM) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/entrypoints/llm + - tests/entrypoints/offline_mode + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py + - pytest -v -s entrypoints/llm/test_generate.py + - pytest -v -s entrypoints/offline_mode + +- label: Entrypoints Integration (Pooling) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/entrypoints/pooling + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/pooling + +- label: Entrypoints Integration (Responses API) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/entrypoints/openai/responses + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/responses + +- label: Entrypoints Unit Tests # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/entrypoints + - tests/entrypoints/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s entrypoints/openai/tool_parsers + - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling + +- label: OpenAI API correctness # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/ + - vllm/entrypoints/openai/ + - vllm/model_executor/models/whisper.py + - vllm/model_executor/layers/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/model_executor/model_loader/ + commands: + - bash ../tools/install_torchcodec_rocm.sh || exit 1 + - pytest -s entrypoints/openai/correctness/ + +#----------------------------------------------------------- mi300 · evals -----------------------------------------------------------# + +- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100-MI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + num_gpus: 1 + working_dir: "/vllm-workspace" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/backends/mla/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030 + +- label: LM Eval Small Models # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt + +- label: LM Eval Small Models (MI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt + +- label: GPQA Eval (GPT-OSS) (2xH100-2xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/layers/fused_moe/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - tests/evals/gpt_oss/ + commands: + - uv pip install --system 'gpt-oss[eval]==0.0.5' + - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx942.txt + +- label: LM Eval Small Models (2xB200-2xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt + +- label: DeepSeek V2-Lite Accuracy (4xH100-4xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 num_gpus: 4 optional: true - working_dir: "/vllm-workspace/tests" + working_dir: "/vllm-workspace" source_file_dependencies: - - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/compilation/ - - tests/distributed/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s distributed/test_elastic_ep.py - - -- label: Engine # 11.3m - timeout_in_minutes: 35 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/engine - - tests/test_sequence - - tests/test_config - - tests/test_logger - - tests/test_vllm_port - commands: - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py - - -- label: Engine (1 GPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/engine/ - - tests/v1/engine/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/engine/test_preprocess_error_handling.py - - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py - - -- label: e2e Scheduling (1 GPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/ - - tests/v1/e2e/general/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/general/test_async_scheduling.py - - -- label: e2e Core (1 GPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/ - - tests/v1/e2e/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py - - -- label: Spec Decode Eagle # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/models/ - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness" - - -- label: Spec Decode Speculators + MTP # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - vllm/transformers_utils/configs/speculators/ - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" - - -- label: Spec Decode Ngram + Suffix # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix" - - -- label: Spec Decode Draft Model # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" - - -- label: V1 e2e (2 GPUs) # 7.1m - timeout_in_minutes: 12 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/e2e - commands: - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" - - -- label: V1 e2e (4 GPUs) # 52.6m - timeout_in_minutes: 106 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/e2e - commands: - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy" - - -- label: V1 e2e (4xH100-4xMI325) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - optional: true - source_file_dependencies: - - vllm/v1/attention/backends/utils.py - - vllm/v1/worker/gpu_model_runner.py - - tests/v1/e2e/test_hybrid_chunked_prefill.py - commands: - - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py - - -- label: V1 Spec Decode # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/spec_decode - commands: - - pytest -v -s -m 'not slow_test' v1/spec_decode - - -- label: V1 Sample + Logits # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/sample - - tests/v1/logits_processors - - tests/v1/test_oracle.py - - tests/v1/test_request.py - - tests/v1/test_outputs.py - commands: - - pytest -v -s v1/sample - - pytest -v -s v1/logits_processors - - pytest -v -s v1/test_oracle.py - - pytest -v -s v1/test_request.py - - pytest -v -s v1/test_outputs.py - - -- label: V1 Core + KV + Metrics # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/core - - tests/v1/executor - - tests/v1/kv_offload - - tests/v1/worker - - tests/v1/kv_connector/unit - - tests/v1/metrics - - tests/entrypoints/openai/correctness/test_lmeval.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - pytest -v -s -m 'not cpu_test' v1/core - - pytest -v -s v1/executor - - pytest -v -s v1/kv_offload - - pytest -v -s v1/worker - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'not cpu_test' v1/metrics - - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api - # - export HSA_NO_SCRATCH_RECLAIM=1 - - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - - - -- label: Acceptance Length Test (Large Models) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/model_executor/models/mlp_speculator.py - - tests/v1/spec_decode/test_acceptance_length.py - - vllm/platforms/rocm.py - commands: - - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - - pytest -v -s v1/spec_decode/test_acceptance_length.py - - -- label: V1 attention (H100-MI325) # 14.5m - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/config/attention.py - - vllm/model_executor/layers/attention - - vllm/v1/attention - - tests/v1/attention + - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/backends/mla/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ - vllm/_aiter_ops.py - - vllm/envs.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/attention + - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010 - -- label: Batch Invariance (H100-MI325) # 5.2m - timeout_in_minutes: 12 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: LM Eval Large Models (4xA100-4xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 optional: true - working_dir: "/vllm-workspace/tests" + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" source_file_dependencies: - - vllm/v1/attention - - vllm/model_executor/layers - - tests/v1/determinism/ + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pip install pytest-timeout pytest-forked - - pytest -v -s v1/determinism/test_batch_invariance.py - - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 - -- label: V1 others (CPU) # 10.4m - timeout_in_minutes: 28 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - no_gpu: true +- label: Qwen3-30B-A3B-FP8-block Accuracy (4xH100-4xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 optional: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 + +- label: Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/spec_decode/ + - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040 + +- label: LM Eval Large Models (8xH200-8xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_8 + optional: true + num_gpus: 8 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/v1 + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/layers/layernorm.py + - csrc/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - tests/evals/ commands: - - pytest -v -s -m 'cpu_test' v1/core - - pytest -v -s v1/structured_output - - pytest -v -s v1/test_serial_utils.py - - pytest -v -s -m 'cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'cpu_test' v1/metrics + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt +#--------------------------------------------------------- mi300 · examples ----------------------------------------------------------# -- label: Examples # 24.5m - timeout_in_minutes: 55 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Examples # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/examples" source_file_dependencies: @@ -1962,54 +1660,12 @@ steps: - python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048 - python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536 +#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------# -- label: Platform Tests (CUDA) # 5.0m - timeout_in_minutes: 9 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/cuda - commands: - - pytest -v -s cuda/test_cuda_context.py - - pytest -v -s cuda/test_platform_no_cuda_init.py - - -- label: PyTorch Compilation Passes Unit Tests # TBD +- label: Kernels Attention Test %N # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/compile/passes - commands: - - pytest -s -v compile/passes --ignore compile/passes/distributed - - -- label: Kernels Core Operation Test # 26.8m - timeout_in_minutes: 38 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - tests/kernels/core - - tests/kernels/test_top_k_per_row.py - - tests/kernels/test_concat_mla_q.py - - vllm/model_executor/layers/rotary_embedding/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s kernels/core kernels/test_top_k_per_row.py - - -- label: Kernels Attention Test %N # 17.7m - timeout_in_minutes: 28 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 parallelism: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2023,29 +1679,26 @@ steps: commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - -- label: Kernels Quantization Test %N # 15.2m - timeout_in_minutes: 24 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - parallelism: 2 +- label: Kernels Core Operation Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/quantization/ - - vllm/model_executor/layers/quantization - - tests/kernels/quantization - - tests/kernels/quantization/test_rocm_skinny_gemms.py + - csrc/ + - tests/kernels/core + - tests/kernels/test_top_k_per_row.py + - tests/kernels/test_concat_mla_q.py + - vllm/model_executor/layers/rotary_embedding/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py - - vllm/model_executor/kernels/ commands: - - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py - label: Kernels MoE Test %N # TBD - timeout_in_minutes: 19 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 parallelism: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2062,11 +1715,27 @@ steps: - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - -- label: Kernels FP8 MoE Test # TBD +- label: Kernels Quantization Test %N # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + parallelism: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/quantization/ + - vllm/model_executor/layers/quantization + - tests/kernels/quantization + - tests/kernels/quantization/test_rocm_skinny_gemms.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/model_executor/kernels/ + commands: + - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +- label: Kernels FP8 MoE Test (2xH100-2xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2080,65 +1749,45 @@ steps: commands: - pytest -v -s kernels/moe/test_deepep_moe.py +#----------------------------------------------------------- mi300 · lora ------------------------------------------------------------# -- label: ROCm AITER Ops Test # TBD +- label: LoRA TP (Distributed) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/_aiter_ops.py - - vllm/envs.py - - vllm/platforms/rocm.py - - tests/rocm/aiter/ - - vllm/v1/attention/backends/mla/rocm_aiter_mla.py - - vllm/v1/attention/selector.py - commands: - - pytest -v -s rocm/aiter/ - - -- label: Benchmarks # 8.2m - timeout_in_minutes: 20 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/.buildkite" - source_file_dependencies: - - benchmarks/ + - vllm/lora + - tests/lora - vllm/platforms/rocm.py commands: - - bash scripts/run-benchmarks.sh + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + - pytest -v -s -x lora/test_chatglm3_tp.py + - pytest -v -s -x lora/test_llama_tp.py + - pytest -v -s -x lora/test_llm_with_multi_loras.py + - pytest -v -s -x lora/test_olmoe_tp.py + - pytest -v -s -x lora/test_gptoss_tp.py + - pytest -v -s -x lora/test_qwen35_densemodel_lora.py +#---------------------------------------------------------- mi300 · models -----------------------------------------------------------# -- label: Quantization # 36.1m - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Language Models Test (Extended Pooling) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - tests/quantization + - vllm/ + - tests/models/language/pooling commands: + - pytest -v -s models/language/pooling -m 'not core_model' - # temporary install here since we need nightly, will move to requirements/test.in - # after torchao 0.12 release, and pin a working version of torchao nightly here - - # since torchao nightly is only compatible with torch nightly currently - # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now - # we can only upgrade after this is resolved - # TODO(jerryzh168): resolve the above comment - - uv pip install --system torchao==0.17.0 - - uv pip install --system conch-triton-kernels - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py - - -- label: Language Models Tests (Standard) # 22.8m - timeout_in_minutes: 38 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Language Models Tests (Standard) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true torch_nightly: true working_dir: "/vllm-workspace/tests" @@ -2149,58 +1798,51 @@ steps: - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and (not slow_test)' - -- label: Language Models Tests (Hybrid) %N # 34.9m - timeout_in_minutes: 55 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - torch_nightly: true - parallelism: 2 +- label: Multi-Modal Models (Extended Generation 1) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/language/generation + - tests/models/multimodal/generation + - tests/models/multimodal/test_mapping.py commands: - - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2' - - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py + - pytest -v -s models/multimodal/test_mapping.py - -- label: Language Models Test (Extended Generation) # 32.2m - timeout_in_minutes: 55 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/language/generation - commands: - - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2' - - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' - - -- label: Multi-Modal Processor # 1h 42m - timeout_in_minutes: 138 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Multi-Modal Models (Extended Generation 2) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/multimodal - - tests/models/registry.py + - tests/models/multimodal/generation commands: - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/processing/test_tensor_schema.py + - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' +- label: Multi-Modal Models (Extended Generation 3) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal/generation + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' - label: "Multi-Modal Models (Standard) 1: qwen2" # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 torch_nightly: true - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -2210,26 +1852,10 @@ steps: - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model - -- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" - - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model - - - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 torch_nightly: true optional: true working_dir: "/vllm-workspace/tests" @@ -2242,29 +1868,11 @@ steps: - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model - - label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 torch_nightly: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing - - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model - - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model - - -- label: Multi-Modal Models (Extended Generation 1) # 1h 2m - timeout_in_minutes: 106 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -2272,54 +1880,43 @@ steps: - tests/models/multimodal/test_mapping.py commands: - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py - - pytest -v -s models/multimodal/test_mapping.py + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model + - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model - -- label: Multi-Modal Models (Extended Generation 2) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Multi-Modal Processor # 1h 42m + timeout_in_minutes: 138 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/multimodal/generation + - tests/models/multimodal + - tests/models/registry.py commands: - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' + - pytest -v -s models/multimodal/processing/test_tensor_schema.py - -- label: Multi-Modal Models (Extended Generation 3) # TBD +- label: Multi-Modal Processor (CPU) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + no_gpu: true optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/multimodal/generation + - tests/models/multimodal + - tests/models/registry.py commands: - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' + - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py - -- label: Multi-Modal Models (Extended Pooling) # TBD +- label: Quantized Models Test # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/pooling - commands: - - pytest -v -s models/multimodal/pooling -m 'not core_model' - - -- label: Quantized Models Test # 21.4m - timeout_in_minutes: 38 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/model_executor/layers/quantization @@ -2330,11 +1927,10 @@ steps: commands: - pytest -v -s models/quantization - -- label: Transformers Nightly Models # 50.9m - timeout_in_minutes: 102 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Transformers Nightly Models # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/" source_file_dependencies: @@ -2358,38 +1954,319 @@ steps: - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper +#------------------------------------------------------- mi300 · quantization --------------------------------------------------------# -- label: Quantized MoE Test (B200-MI325) # TBD +- label: Quantization # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/" + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - tests/quantization/test_gfx3xx_moe.py - - vllm/model_executor/models/deepseek_v2.py - - vllm/model_executor/models/gpt_oss.py - - vllm/model_executor/models/llama4.py - - vllm/model_executor/layers/fused_moe - - vllm/model_executor/layers/quantization/compressed_tensors - - vllm/model_executor/layers/quantization/modelopt.py - - vllm/model_executor/layers/quantization/mxfp4.py - - vllm/v1/attention/backends/triton_attn.py - - vllm/v1/attention/backends/rocm_attn.py - - vllm/v1/attention/backends/rocm_aiter_fa.py - - vllm/v1/attention/backends/mla/ - - vllm/v1/attention/selector.py - - vllm/model_executor/layers/layernorm.py + - csrc/ + - vllm/model_executor/layers/quantization - vllm/_aiter_ops.py - vllm/platforms/rocm.py - - vllm/model_executor/model_loader/ + - tests/quantization commands: - - pytest -s -v tests/quantization/test_gfx3xx_moe.py + # temporary install here since we need nightly, will move to requirements/test.in + # after torchao 0.12 release, and pin a working version of torchao nightly here -- label: Distributed DP Tests (2 GPUs) # 56.1m - timeout_in_minutes: 102 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 + # since torchao nightly is only compatible with torch nightly currently + # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now + # we can only upgrade after this is resolved + # TODO(jerryzh168): resolve the above comment + - uv pip install --system torchao==0.17.0 + - uv pip install --system conch-triton-kernels + - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py + +#----------------------------------------------------------- mi300 · rocm ------------------------------------------------------------# + +- label: ROCm AITER Ops Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/_aiter_ops.py + - vllm/envs.py + - vllm/platforms/rocm.py + - tests/rocm/aiter/ + - vllm/v1/attention/backends/mla/rocm_aiter_mla.py + - vllm/v1/attention/selector.py + commands: + - pytest -v -s rocm/aiter/ + +#--------------------------------------------------------- mi300 · samplers ----------------------------------------------------------# + +- label: Samplers Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/model_executor/layers + - vllm/sampling_metadata.py + - vllm/v1/sample/ + - vllm/beam_search.py + - tests/samplers + - tests/conftest.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s samplers + +#------------------------------------------------------------ mi300 · misc ------------------------------------------------------------# + +- label: Python-only Installation # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - tests/standalone_tests/python_only_compile.sh + - setup.py + - vllm/platforms/rocm.py + commands: + - bash standalone_tests/python_only_compile.sh + +- label: Regression # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/test_regression + commands: + - pip install modelscope + - pytest -v -s test_regression.py + +#--------------------------------------------------------- mi300 · ray_compat ---------------------------------------------------------# + +- label: Ray Dependency Compatibility Check # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/" + source_file_dependencies: + - requirements/ + - setup.py + - vllm/platforms/rocm.py + commands: + - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh + +#------------------------------------------------------------ mi300 · v1 -------------------------------------------------------------# + +- label: Acceptance Length Test (Large Models) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/model_executor/models/mlp_speculator.py + - tests/v1/spec_decode/test_acceptance_length.py + - vllm/platforms/rocm.py + commands: + - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 + - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test + +- label: e2e Core (1 GPU) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/ + - tests/v1/e2e/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py + +- label: e2e Scheduling (1 GPU) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/ + - tests/v1/e2e/general/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/general/test_async_scheduling.py + +- label: Engine (1 GPU) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/engine/ + - tests/v1/engine/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/engine/test_preprocess_error_handling.py + - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py + +- label: Spec Decode Draft Model # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" + +- label: Spec Decode Eagle # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness" + +- label: Spec Decode Ngram + Suffix # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix" + +- label: Spec Decode Speculators + MTP # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - vllm/transformers_utils/configs/speculators/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" + +- label: V1 attention (H100-MI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/config/attention.py + - vllm/model_executor/layers/attention + - vllm/v1/attention + - tests/v1/attention + - vllm/_aiter_ops.py + - vllm/envs.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/attention + +- label: V1 Core + KV + Metrics # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - pytest -v -s -m 'not cpu_test' v1/core + - pytest -v -s v1/executor + - pytest -v -s v1/kv_offload + - pytest -v -s v1/worker + - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'not cpu_test' v1/metrics + - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api + # - export HSA_NO_SCRATCH_RECLAIM=1 + - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine + +- label: V1 others (CPU) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + no_gpu: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1 + commands: + - pytest -v -s -m 'cpu_test' v1/core + - pytest -v -s v1/structured_output + - pytest -v -s v1/test_serial_utils.py + - pytest -v -s -m 'cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'cpu_test' v1/metrics + +- label: V1 Sample + Logits # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/sample + - tests/v1/logits_processors + - tests/v1/test_oracle.py + - tests/v1/test_request.py + - tests/v1/test_outputs.py + commands: + - pytest -v -s v1/sample + - pytest -v -s v1/logits_processors + - pytest -v -s v1/test_oracle.py + - pytest -v -s v1/test_request.py + - pytest -v -s v1/test_outputs.py + +- label: Distributed DP Tests (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 num_gpus: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2409,9 +2286,196 @@ steps: - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py +- label: Distributed Tests (2xH100-2xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + working_dir: "/vllm-workspace/" + source_file_dependencies: + - vllm/distributed/ + - vllm/v1/distributed/ + - vllm/model_executor/layers/fused_moe/ + - tests/v1/distributed/test_dbo.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export TORCH_NCCL_BLOCKING_WAIT=1 + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py + - VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput + - pytest -v -s tests/v1/distributed/test_dbo.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py + - pytest -v -s tests/distributed/test_packed_tensor.py -- label: Distributed Compile + RPC Tests (2 GPUs) # 56.1m - timeout_in_minutes: 102 +- label: Metrics, Tracing (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/tracing + commands: + - "pip install \ + 'opentelemetry-sdk>=1.26.0' \ + 'opentelemetry-api>=1.26.0' \ + 'opentelemetry-exporter-otlp>=1.26.0' \ + 'opentelemetry-semantic-conventions-ai>=0.4.1'" + - pytest -v -s v1/tracing + +- label: V1 e2e (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/e2e + commands: + - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" + +- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - CROSS_LAYERS_BLOCKS=True ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: Distributed DP Tests (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/ + - tests/v1/distributed + - tests/v1/engine/test_engine_core_client.py + - tests/distributed/test_utils + - vllm/platforms/rocm.py + commands: + - export TORCH_NCCL_BLOCKING_WAIT=1 + - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py + - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py + - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py + - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp + - pytest -v -s distributed/test_utils.py + +- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - HYBRID_SSM=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: V1 e2e (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/e2e + commands: + - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy" + +- label: V1 e2e (4xH100-4xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + optional: true + source_file_dependencies: + - vllm/v1/attention/backends/utils.py + - vllm/v1/worker/gpu_model_runner.py + - tests/v1/e2e/test_hybrid_chunked_prefill.py + commands: + - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py + +#------------------------------------------------------ mi300 · weight_loading -------------------------------------------------------# + +- label: Weight Loading Multiple GPU # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/weight_loading + commands: + - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt + +- label: Weight Loading Multiple GPU - Large Models # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/weight_loading + commands: + - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt + +######################################################################################################################################### +# # +# MI325 (gfx942) tests # +# # +######################################################################################################################################### + +#---------------------------------------------------------- mi325 · compile ----------------------------------------------------------# + +- label: Distributed Compile + RPC Tests (2 GPUs) # TBD + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_2 num_gpus: 2 @@ -2434,9 +2498,10 @@ steps: - pytest -v -s ./compile/fullgraph/test_basic_correctness.py - pytest -v -s ./compile/test_wrapper.py +#-------------------------------------------------------- mi325 · distributed --------------------------------------------------------# -- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # 56.1m - timeout_in_minutes: 102 +- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_2 num_gpus: 2 @@ -2459,266 +2524,50 @@ steps: - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - pytest -v -s v1/worker/test_worker_memory_snapshot.py - -- label: Distributed Model Tests (2 GPUs) # 19.3m - timeout_in_minutes: 38 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/model_executor/model_loader/sharded_state_loader.py - - vllm/model_executor/models/ - - vllm/model_executor/layers/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - tests/basic_correctness/ - - tests/model_executor/model_loader/test_sharded_state_loader.py - - tests/models/ - commands: - - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' - - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' - - pytest models/language -v -s -m 'distributed(num_gpus=2)' - - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py - - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' - - -- label: LoRA TP (Distributed) # 9.8m - timeout_in_minutes: 18 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/lora - - tests/lora - - vllm/platforms/rocm.py - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True - - pytest -v -s -x lora/test_chatglm3_tp.py - - pytest -v -s -x lora/test_llama_tp.py - - pytest -v -s -x lora/test_llm_with_multi_loras.py - - pytest -v -s -x lora/test_olmoe_tp.py - - pytest -v -s -x lora/test_gptoss_tp.py - - pytest -v -s -x lora/test_qwen35_densemodel_lora.py - - -- label: Weight Loading Multiple GPU # 7.5m - timeout_in_minutes: 14 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - num_gpus: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/weight_loading - commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt - - -- label: Weight Loading Multiple GPU - Large Models # 12.6m - timeout_in_minutes: 26 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/weight_loading - commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt - - -- label: Ray Dependency Compatibility Check # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/" - source_file_dependencies: - - requirements/ - - setup.py - - vllm/platforms/rocm.py - commands: - - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh - - -- label: Distributed NixlConnector PD accuracy (4 GPUs) # 27.4m - timeout_in_minutes: 44 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - - -- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD +- label: Distributed Compile + Comm (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_4 num_gpus: 4 - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - CROSS_LAYERS_BLOCKS=True ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - - -- label: Distributed Tests (4 GPUs)(A100-MI325) # 20.9m - timeout_in_minutes: 37 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - pytest -v -s distributed/test_custom_all_reduce.py - - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py - - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - - pytest -v -s -x lora/test_mixtral.py - - -- label: Distributed Tests (2 GPUs)(H100-MI325) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - num_gpus: 2 - working_dir: "/vllm-workspace/" - source_file_dependencies: - vllm/distributed/ - - vllm/v1/distributed/ - - vllm/model_executor/layers/fused_moe/ - - tests/v1/distributed/test_dbo.py - - vllm/_aiter_ops.py + - tests/distributed/test_pynccl + - tests/distributed/test_events + - tests/compile/fullgraph/test_basic_correctness.py + - tests/distributed/test_symm_mem_allreduce.py + - tests/distributed/test_multiproc_executor.py - vllm/platforms/rocm.py commands: - export TORCH_NCCL_BLOCKING_WAIT=1 - - VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - - pytest -v -s tests/v1/distributed/test_dbo.py + - pytest -v -s compile/fullgraph/test_basic_correctness.py + - pytest -v -s distributed/test_pynccl.py + - pytest -v -s distributed/test_events.py + - pytest -v -s distributed/test_symm_mem_allreduce.py + - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node +#---------------------------------------------------------- mi325 · engine -----------------------------------------------------------# -- label: Distributed Compile Unit Tests (2xH100-2xMI325) # 14.3m - timeout_in_minutes: 32 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - num_gpus: 2 - working_dir: "/vllm-workspace/" - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers - - tests/compile/passes/distributed/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py - - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py - # TODO: this test is not supported on ROCm, there are aiter kernels for this. - # - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py - # - pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm - # - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'" - - -- label: LM Eval Small Models # 13.3m - timeout_in_minutes: 23 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt - - -- label: LM Eval Small Models (MI325) # TBD +- label: Engine # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_1 - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt - - -- label: LM Eval Small Models (B200-MI325) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/engine + - tests/test_sequence + - tests/test_config + - tests/test_logger + - tests/test_vllm_port commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt + - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py +#----------------------------------------------------------- mi325 · evals -----------------------------------------------------------# -- label: LM Eval Large Models (H200-MI325) # TBD +- label: LM Eval Large Models (4xH100-4xMI325) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_8 - optional: true - num_gpus: 8 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/model_executor/layers/layernorm.py - - csrc/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - tests/evals/ - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt - - -- label: LM Eval Large Models (4 GPUs)(FP8) # 24.8m - timeout_in_minutes: 42 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_4 num_gpus: 4 optional: true @@ -2736,29 +2585,7 @@ steps: - export VLLM_USE_DEEP_GEMM=0 - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4 - -- label: LM Eval Large Models (4 GPUs)(A100-MI325) # 17.3m - timeout_in_minutes: 27 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 - - -- label: ROCm LM Eval Large Models (8 Card) # TBD +- label: ROCm LM Eval Large Models (8 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_8 @@ -2779,264 +2606,122 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8 +#---------------------------------------------------------- mi325 · models -----------------------------------------------------------# -- label: GPQA Eval (GPT-OSS) (H100-MI325) # TBD +- label: Language Models Test (Extended Generation) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 + agent_pool: mi325_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/language/generation + commands: + - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' + - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' + +- label: Language Models Tests (Hybrid) %N # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] + agent_pool: mi325_1 + torch_nightly: true + parallelism: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/language/generation + commands: + - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' + - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + +- label: Multi-Modal Models (Extended Pooling) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] + agent_pool: mi325_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal/pooling + commands: + - pytest -v -s models/multimodal/pooling -m 'not core_model' + +- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] + agent_pool: mi325_1 + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" + - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model + +#------------------------------------------------------------ mi325 · v1 -------------------------------------------------------------# + +- label: V1 Spec Decode # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] + agent_pool: mi325_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/spec_decode + commands: + - pytest -v -s -m 'not slow_test' v1/spec_decode + +######################################################################################################################################### +# # +# MI355 (gfx950) tests # +# # +######################################################################################################################################### + +#-------------------------------------------------------- mi355 · benchmarks ---------------------------------------------------------# + +- label: Attention Benchmarks Smoke Test (B200-MI355) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_2 + num_gpus: 2 + working_dir: "/vllm-workspace/" + source_file_dependencies: + - benchmarks/attention_benchmarks/ + - vllm/v1/attention/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1 + +#-------------------------------------------------------- mi355 · distributed --------------------------------------------------------# + +- label: Distributed Tests (2xH100-2xMI355) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_2 num_gpus: 2 optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/model_executor/layers/fused_moe/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - tests/evals/gpt_oss/ - commands: - - uv pip install --system 'gpt-oss[eval]==0.0.5' - - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx942.txt - - -- label: DeepSeek V2-Lite Accuracy # 6.7m - timeout_in_minutes: 12 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/distributed/eplb - - vllm/model_executor/layers/fused_moe/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/backends/mla/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010 - - -- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100-MI325) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - num_gpus: 1 - working_dir: "/vllm-workspace" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/fused_moe/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/backends/mla/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030 - - -- label: Qwen3-30B-A3B-FP8-block Accuracy # 6.4m - timeout_in_minutes: 11 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - optional: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/quantization/ - - vllm/distributed/eplb - - vllm/model_executor/layers/fused_moe/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 - - -- label: Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy # 10.9m - timeout_in_minutes: 22 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/spec_decode/ - - vllm/distributed/eplb - - vllm/model_executor/layers/fused_moe/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040 - -##### .buildkite/test_areas/compile.yaml ##### -# Slowly setting up the tests so that it is also easier for the -# CI team to review and upstream to the pipelinev2. -# The following tests are important for vLLM IR Ops refactoring, -# which affects fusion passes on ROCm. So we have to -# enable them as as soon as possible. - -## TODO: Enable the test in this group -# # corresponds to .buildkite/test_areas/compile.yaml -# - label: Fusion and Compile Unit Tests (2xB200-2xMI325) # TBD -# timeout_in_minutes: 180 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325, tj] -# agent_pool: mi325_1 # changed to 1 GPU until the fusion all reduce is enabled then only revert back to 2 GPUs -# num_gpus: 1 -# working_dir: "/vllm-workspace/" -# source_file_dependencies: -# - csrc/quantization/fp4/ -# - vllm/model_executor/layers/quantization/ -# - vllm/model_executor/layers/layernorm.py -# - vllm/model_executor/layers/activation.py -# - vllm/model_executor/layers/attention/attention.py -# - vllm/v1/attention/backends/flashinfer.py -# - vllm/compilation/ # TODO(luka) limit to vllm/compilation/passes -# - tests/compile/test_fusion_attn.py -# - tests/compile/test_silu_mul_quant_fusion.py -# - tests/compile/distributed/test_fusion_all_reduce.py -# - tests/compile/fullgraph/test_full_graph.py -# commands: -# - rocm-smi -# # we run all backend tests on ROCm -# # These two tests are covered in "PyTorch Compilation Passes Unit Tests" -# # - "pytest -v -s tests/compile/passes/test_fusion_attn.py" -# # - "pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py" -# # TODO: this test is not supported on ROCm, there are aiter kernels for this. -# # - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py -# # TODO: find out more details -# # - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile - - -- label: Fusion E2E Quick (H100-MI325) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - num_gpus: 1 working_dir: "/vllm-workspace/" source_file_dependencies: - - csrc/quantization/ - - vllm/model_executor/ - - vllm/v1/attention/ - - vllm/compilation/ - - tests/compile/fusions_e2e/ + - vllm/distributed/ + - vllm/v1/distributed/ + - vllm/model_executor/layers/fused_moe/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - tests/distributed/test_context_parallel.py + - tests/v1/distributed/test_dbo.py + - examples/offline_inference/data_parallel.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - rocm-smi - # Run all models and attn backends but only Inductor partition and native custom ops - - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and not +rms_norm and not +quant_fp8'" - # Different from CUDA, Qwen requires +rms_norm and +quant_fp8 as rms+quant fusion is only supported on AITER - - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and +rms_norm and +quant_fp8 and qwen3'" - - -- label: Fusion E2E Config Sweep (H100-MI325) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - num_gpus: 1 - working_dir: "/vllm-workspace/" - source_file_dependencies: - - csrc/quantization/ - - vllm/compilation/ - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/attention/attention.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/fusions_e2e/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - rocm-smi - - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3" - -## There are no ops on ROCm for these tests. -## The test still passes but the logs are not useful. -## fused ops just call torch.ops.symm_mem which -## exists in ROCm even though they don't work -# - label: AsyncTP Correctness Tests (2xH100-2xMI325) -# - label: Fusion E2E TP2 Quick (H100-MI325) -# - label: Fusion E2E TP2 AsyncTP Config Sweep (H100-MI325) -# - label: Fusion E2E TP2 (B200-MI325) -# - label: Sequence Parallel Correctness Tests (2xH100-2xMI325) - - -##################################################################################################################################### -# # -# gfx950 # -# # -##################################################################################################################################### - -- label: Entrypoints Integration (API Server openai - Part 1) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/openai - - tests/entrypoints/test_chat_utils - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py - - -- label: Entrypoints Integration (API Server openai - Part 2) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/openai - - tests/entrypoints/test_chat_utils - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - - pytest -v -s entrypoints/openai/speech_to_text/ - - pytest -v -s entrypoints/test_chat_utils.py - - -- label: Entrypoints Integration (API Server openai - Part 3) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/openai - - tests/entrypoints/test_chat_utils - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses + - export TORCH_NCCL_BLOCKING_WAIT=1 + - pytest -v -s tests/distributed/test_context_parallel.py + - pytest -v -s tests/v1/distributed/test_dbo.py +#-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------# - label: Entrypoints Integration (API Server 2) # TBD timeout_in_minutes: 180 @@ -3057,6 +2742,52 @@ steps: - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc - pytest -v -s tool_use +- label: Entrypoints Integration (API Server openai - Part 1) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/entrypoints/openai + - tests/entrypoints/test_chat_utils + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py + +- label: Entrypoints Integration (API Server openai - Part 2) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/entrypoints/openai + - tests/entrypoints/test_chat_utils + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py + - pytest -v -s entrypoints/openai/speech_to_text/ + - pytest -v -s entrypoints/test_chat_utils.py + +- label: Entrypoints Integration (API Server openai - Part 3) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/entrypoints/openai + - tests/entrypoints/test_chat_utils + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py - label: Entrypoints Integration (Pooling) # TBD timeout_in_minutes: 180 @@ -3072,95 +2803,112 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/pooling +#----------------------------------------------------------- mi355 · evals -----------------------------------------------------------# -- label: Regression # TBD +- label: GPQA Eval (GPT-OSS) (2xB200-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 + agent_pool: mi355_2 + num_gpus: 2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/test_regression - commands: - - pip install modelscope - - pytest -v -s test_regression.py - - -- label: V1 Spec Decode # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/spec_decode - commands: - - pytest -v -s -m 'not slow_test' v1/spec_decode - - -- label: V1 Sample + Logits # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/sample - - tests/v1/logits_processors - - tests/v1/test_oracle.py - - tests/v1/test_request.py - - tests/v1/test_outputs.py - commands: - - pytest -v -s v1/sample - - pytest -v -s v1/logits_processors - - pytest -v -s v1/test_oracle.py - - pytest -v -s v1/test_request.py - - pytest -v -s v1/test_outputs.py - - -- label: V1 Core + KV + Metrics # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/core - - tests/v1/executor - - tests/v1/kv_offload - - tests/v1/worker - - tests/v1/kv_connector/unit - - tests/v1/metrics - - tests/entrypoints/openai/correctness/test_lmeval.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - pytest -v -s -m 'not cpu_test' v1/core - - pytest -v -s v1/executor - - pytest -v -s v1/kv_offload - - pytest -v -s v1/worker - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'not cpu_test' v1/metrics - - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api - - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - - -- label: V1 attention (B200-MI355) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/config/attention.py - - vllm/model_executor/layers/attention - - vllm/v1/attention - - tests/v1/attention + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/layers/fused_moe/ + - tests/evals/gpt_oss/ - vllm/_aiter_ops.py - - vllm/envs.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/attention + - uv pip install --system 'gpt-oss[eval]==0.0.5' + - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt +- label: LM Eval Qwen3-5 Models (B200-MI355) # TBD + timeout_in_minutes: 120 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py + - vllm/transformers_utils/configs/qwen3_5.py + - vllm/transformers_utils/configs/qwen3_5_moe.py + - vllm/model_executor/models/qwen.py + - vllm/model_executor/models/qwen2.py + - vllm/model_executor/models/qwen3.py + - vllm/model_executor/models/qwen3_next.py + - vllm/model_executor/models/qwen3_next_mtp.py + - vllm/model_executor/layers/fla/ops/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt + +- label: LM Eval Small Models (2xB200-2xMI355) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt + +- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_2 + num_gpus: 2 + working_dir: "/vllm-workspace" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/quantization/ + - vllm/model_executor/layers/fused_moe/ + - vllm/distributed/eplb + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 + +- label: LM Eval Large Models (4xH100-4xMI355) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_DEEP_GEMM=0 + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4 + +#--------------------------------------------------------- mi355 · examples ----------------------------------------------------------# - label: Examples # TBD timeout_in_minutes: 180 @@ -3195,276 +2943,7 @@ steps: - python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048 - python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536 - -- label: Kernels Attention Test %N # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - parallelism: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/attention/ - - vllm/v1/attention - - vllm/model_executor/layers/attention - - tests/kernels/attention - - vllm/_aiter_ops.py - - vllm/envs.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - -- label: Kernels Quantization Test %N # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - parallelism: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/quantization/ - - vllm/model_executor/layers/quantization - - tests/kernels/quantization - - tests/kernels/quantization/test_rocm_skinny_gemms.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/model_executor/kernels/ - commands: - - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - -- label: Kernels MoE Test %N # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - parallelism: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ - - csrc/moe/ - - tests/kernels/moe - - vllm/model_executor/layers/fused_moe/ - - vllm/distributed/device_communicators/ - - vllm/envs.py - - vllm/config - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - -- label: Kernels FP8 MoE Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/moe/ - - csrc/quantization/w8a8/cutlass/moe/ - - vllm/model_executor/layers/fused_moe/ - - tests/kernels/moe/test_deepep_moe.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/envs.py - commands: - - pytest -v -s kernels/moe/test_deepep_moe.py - - -- label: Quantization # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - tests/quantization - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - uv pip install --system torchao==0.17.0 - - uv pip install --system conch-triton-kernels - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py - - -- label: Language Models Tests (Standard) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/language - commands: - - pip freeze | grep -E 'torch' - - pytest -v -s models/language -m 'core_model and (not slow_test)' - - -- label: Language Models Test (Extended Generation) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/language/generation - commands: - - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2' - - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' - - -- label: Language Models Test (Extended Pooling) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/language/pooling - commands: - - pytest -v -s models/language/pooling -m 'not core_model' - - -- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - torch_nightly: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" - - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model - - -- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - torch_nightly: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" - - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model - - -- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - torch_nightly: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - - tests/models/multimodal/test_mapping.py - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" - - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model - - -- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - torch_nightly: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing - - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model - - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model - - -- label: Multi-Modal Models (Extended Generation 1) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - - tests/models/multimodal/test_mapping.py - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py - - pytest -v -s models/multimodal/test_mapping.py - - -- label: Multi-Modal Models (Extended Generation 2) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' - - -- label: Multi-Modal Models (Extended Generation 3) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' - - -- label: Multi-Modal Models (Extended Pooling) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/pooling - commands: - - pytest -v -s models/multimodal/pooling -m 'not core_model' - - -- label: Quantized Models Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/model_executor/layers/quantization - - tests/models/quantization - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/model_executor/model_loader/ - commands: - - pytest -v -s models/quantization - +#---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------# - label: Kernels (B200-MI355) # TBD timeout_in_minutes: 180 @@ -3490,79 +2969,345 @@ steps: - python3 examples/basic/offline_inference/chat.py - pytest -v -s tests/kernels/attention/test_attention_selector.py - -- label: Weight Loading Multiple GPU # TBD +- label: Kernels Attention Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - num_gpus: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/weight_loading - commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt - - -- label: Weight Loading Multiple GPU - Large Models # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - working_dir: "/vllm-workspace/tests" - num_gpus: 2 + agent_pool: mi355_1 + parallelism: 2 optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/attention/ + - vllm/v1/attention + - vllm/model_executor/layers/attention + - tests/kernels/attention + - vllm/_aiter_ops.py + - vllm/envs.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +- label: Kernels MoE Test %N # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + parallelism: 4 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/quantization/cutlass_w8a8/moe/ + - csrc/moe/ + - tests/kernels/moe + - vllm/model_executor/layers/fused_moe/ + - vllm/distributed/device_communicators/ + - vllm/envs.py + - vllm/config + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +- label: Kernels Quantization Test %N # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + parallelism: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/quantization/ + - vllm/model_executor/layers/quantization + - tests/kernels/quantization + - tests/kernels/quantization/test_rocm_skinny_gemms.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/model_executor/kernels/ + commands: + - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +- label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/moe/ + - csrc/quantization/w8a8/cutlass/moe/ + - vllm/model_executor/layers/fused_moe/ + - tests/kernels/moe/test_deepep_moe.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/envs.py + commands: + - pytest -v -s kernels/moe/test_deepep_moe.py + +#---------------------------------------------------------- mi355 · models -----------------------------------------------------------# + +- label: Language Models Test (Extended Generation) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/weight_loading + - tests/models/language/generation commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt + - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' + - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' +- label: Language Models Test (Extended Pooling) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/language/pooling + commands: + - pytest -v -s models/language/pooling -m 'not core_model' -- label: Ray Dependency Compatibility Check # TBD +- label: Language Models Test (PPL) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py + - vllm/transformers_utils/configs/qwen3_5.py + - vllm/transformers_utils/configs/qwen3_5_moe.py + - vllm/model_executor/models/qwen.py + - vllm/model_executor/models/qwen2.py + - vllm/model_executor/models/qwen3.py + - vllm/model_executor/models/qwen3_next.py + - vllm/model_executor/models/qwen3_next_mtp.py + - vllm/model_executor/layers/fla/ops/ + - vllm/_aiter_ops.py + - vllm/v1/attention/backends/triton_attn.py + - vllm/v1/attention/backends/rocm_attn.py + - vllm/v1/attention/backends/rocm_aiter_unified_attn.py + - vllm/v1/attention/backends/rocm_aiter_fa.py + - vllm/v1/attention/backends/flex_attention.py + - vllm/v1/attention/ops/ + - vllm/platforms/rocm.py + - tests/models/language/generation_ppl_test + commands: + - pytest -v -s models/language/generation_ppl_test + +- label: Language Models Tests (Standard) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/language + commands: + - pip freeze | grep -E 'torch' + - pytest -v -s models/language -m 'core_model and (not slow_test)' + +- label: Multi-Modal Models (Extended Generation 1) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 optional: true - working_dir: "/" + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - requirements/ - - setup.py - - vllm/platforms/rocm.py + - vllm/ + - tests/models/multimodal/generation + - tests/models/multimodal/test_mapping.py commands: - - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py + - pytest -v -s models/multimodal/test_mapping.py - -- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD +- label: Multi-Modal Models (Extended Generation 3) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_4 - num_gpus: 4 + agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py + - vllm/ + - tests/models/multimodal/generation commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' - -- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD +- label: Multi-Modal Models (Extended Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_4 - num_gpus: 4 + agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ + - vllm/ + - tests/models/multimodal/pooling + commands: + - pytest -v -s models/multimodal/pooling -m 'not core_model' + +- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + torch_nightly: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" + - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model + +- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + torch_nightly: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal/generation + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model + - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model + +- label: Quantized Models Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/model_executor/layers/quantization + - tests/models/quantization + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/model_executor/model_loader/ + commands: + - pytest -v -s models/quantization + +#------------------------------------------------------- mi355 · quantization --------------------------------------------------------# + +- label: Quantization # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - tests/quantization + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + - uv pip install --system torchao==0.17.0 + - uv pip install --system conch-triton-kernels + - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py +# - label: Quantized MoE Test (B200-MI355) # TBD +# timeout_in_minutes: 180 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# agent_pool: mi355_1 +# working_dir: "/vllm-workspace/" +# source_file_dependencies: +# - tests/quantization/test_gfx950_moe.py +# - vllm/model_executor/models/deepseek_v2.py +# - vllm/model_executor/models/gpt_oss.py +# - vllm/model_executor/models/llama4.py +# - vllm/model_executor/layers/fused_moe +# - vllm/model_executor/layers/quantization/compressed_tensors +# - vllm/model_executor/layers/quantization/modelopt.py +# - vllm/model_executor/layers/quantization/mxfp4.py +# - vllm/v1/attention/backends/triton_attn.py +# - vllm/v1/attention/backends/rocm_attn.py +# - vllm/v1/attention/backends/rocm_aiter_fa.py +# - vllm/v1/attention/backends/mla/ +# - vllm/v1/attention/selector.py +# - vllm/model_executor/layers/layernorm.py +# - vllm/_aiter_ops.py +# - vllm/platforms/rocm.py +# - vllm/model_executor/model_loader/ +# commands: +# - pytest -s -v tests/quantization/test_gfx950_moe.py + +#------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# + +- label: V1 attention (B200-MI355) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/config/attention.py + - vllm/model_executor/layers/attention + - vllm/v1/attention + - tests/v1/attention + - vllm/_aiter_ops.py + - vllm/envs.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/attention + +- label: V1 Core + KV + Metrics # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - pytest -v -s -m 'not cpu_test' v1/core + - pytest -v -s v1/executor + - pytest -v -s v1/kv_offload + - pytest -v -s v1/worker + - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'not cpu_test' v1/metrics + - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api + - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine + +- label: V1 Sample + Logits # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/sample + - tests/v1/logits_processors + - tests/v1/test_oracle.py + - tests/v1/test_request.py + - tests/v1/test_outputs.py + commands: + - pytest -v -s v1/sample + - pytest -v -s v1/logits_processors + - pytest -v -s v1/test_oracle.py + - pytest -v -s v1/test_request.py + - pytest -v -s v1/test_outputs.py + +- label: V1 Spec Decode # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/spec_decode + commands: + - pytest -v -s -m 'not slow_test' v1/spec_decode - label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD timeout_in_minutes: 180 @@ -3580,172 +3325,74 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh - -- label: Distributed Tests (2 GPUs)(H100-MI355) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/" - source_file_dependencies: - - vllm/distributed/ - - vllm/v1/distributed/ - - vllm/model_executor/layers/fused_moe/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - tests/distributed/test_context_parallel.py - - tests/v1/distributed/test_dbo.py - - examples/offline_inference/data_parallel.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - pytest -v -s tests/distributed/test_context_parallel.py - - pytest -v -s tests/v1/distributed/test_dbo.py - - -- label: Distributed Compile Unit Tests (2xH100-2xMI355) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/" - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers - - tests/compile/passes/distributed/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py - - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py - # TODO: this test is not supported on ROCm, there are aiter kernels for this. - # - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py - # - pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm - # - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'" - - -- label: LM Eval Small Models (B200-MI355) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt - - -- label: LM Eval Large Models (4 GPUs)(FP8) # TBD +- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_4 num_gpus: 4 optional: true - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - export VLLM_USE_DEEP_GEMM=0 - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4 + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - -- label: GPQA Eval (GPT-OSS) (B200-MI355) # TBD +- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx955nightly, amdmi355] - agent_pool: mi355_2 - num_gpus: 2 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_4 + num_gpus: 4 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/model_executor/layers/fused_moe/ - - tests/evals/gpt_oss/ - - vllm/_aiter_ops.py + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system 'gpt-oss[eval]==0.0.5' - - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +#------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------# -- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355) # TBD +- label: Weight Loading Multiple GPU # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_2 num_gpus: 2 - working_dir: "/vllm-workspace" + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/quantization/ - - vllm/model_executor/layers/fused_moe/ - - vllm/distributed/eplb - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/weight_loading commands: - - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 + - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt - -- label: Attention Benchmarks Smoke Test (B200-MI355) # TBD +- label: Weight Loading Multiple GPU - Large Models # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_2 + working_dir: "/vllm-workspace/tests" num_gpus: 2 - working_dir: "/vllm-workspace/" + optional: true source_file_dependencies: - - benchmarks/attention_benchmarks/ - - vllm/v1/attention/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/weight_loading commands: - - python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1 + - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt +#----------------------------------------------------------- mi355 · misc ------------------------------------------------------------# -- label: LM Eval Qwen3-5 Models (B200-MI355) # TBD - timeout_in_minutes: 120 +- label: Regression # TBD + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - num_gpus: 2 + agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/models/qwen3_5.py - - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - - vllm/model_executor/models/qwen.py - - vllm/model_executor/models/qwen2.py - - vllm/model_executor/models/qwen3.py - - vllm/model_executor/models/qwen3_next.py - - vllm/model_executor/models/qwen3_next_mtp.py - - vllm/model_executor/layers/fla/ops/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/test_regression commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt + - pip install modelscope + - pytest -v -s test_regression.py diff --git a/tests/quantization/test_mi3xx_moe.py b/tests/quantization/test_gfx950_moe.py similarity index 58% rename from tests/quantization/test_mi3xx_moe.py rename to tests/quantization/test_gfx950_moe.py index 2f8dfde6847..9cb94086f73 100644 --- a/tests/quantization/test_mi3xx_moe.py +++ b/tests/quantization/test_gfx950_moe.py @@ -2,5 +2,5 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -def test_mi3xx_moe(): - print("TODO: add tests for Mi3xx MoE quantization") +def test_mi355_moe(): + print("TODO: add tests for Mi355 MoE quantization") From 58631d7c3f9984f979e3cd5abf20a61590b36b1f Mon Sep 17 00:00:00 2001 From: nemanjaudovic <152565955+nemanjaudovic@users.noreply.github.com> Date: Mon, 20 Apr 2026 10:58:39 +0200 Subject: [PATCH 135/150] [Bugfix] Fix scaled_mm output narrowing for 3D input tensors (#38093) Signed-off-by: nemanjaudovic --- .../kernels/linear/scaled_mm/pytorch.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py index 2fb6e87413a..ef73be0aafa 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math import torch @@ -13,6 +14,13 @@ from .ScaledMMLinearKernel import ( ) +def _get_num_tokens(output_shape: list) -> int: + # torch._scaled_mm works with 2D tensors, so input tensors are + # flattened if they are 3D. If output_shape is 3D, num_tokens is + # the product of all dims except the last (hidden dim). + return math.prod(output_shape[:-1]) + + class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): """ Base class for FP8 linear kernels using Torch. @@ -78,7 +86,8 @@ class PerTensorTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel): if type(output) is tuple and len(output) == 2: output = output[0] - return torch.narrow(output, 0, 0, output_shape[0]).view(*output_shape) + num_tokens = _get_num_tokens(output_shape) + return torch.narrow(output, 0, 0, num_tokens).view(*output_shape) class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel): @@ -145,7 +154,8 @@ class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel): bias=bias, ) - return torch.narrow(output, 0, 0, output_shape[0]).view(*output_shape) + num_tokens = _get_num_tokens(output_shape) + return torch.narrow(output, 0, 0, num_tokens).view(*output_shape) class ChannelWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel): @@ -206,8 +216,9 @@ class ChannelWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel): output = output[0] # Unpad (undo num_token_padding) - output = torch.narrow(output, 0, 0, output_shape[0]) - x_scale = torch.narrow(As, 0, 0, output_shape[0]) + num_tokens = _get_num_tokens(output_shape) + output = torch.narrow(output, 0, 0, num_tokens) + x_scale = torch.narrow(As, 0, 0, num_tokens) # DQ # C = sw * sx * (X * W) + bias From 2aab9acf48b5f10a4ee0b29f152fc497a46c956d Mon Sep 17 00:00:00 2001 From: Fadi Arafeh <115173828+fadara01@users.noreply.github.com> Date: Mon, 20 Apr 2026 10:21:12 +0100 Subject: [PATCH 136/150] [CPU][BugFix] Fix inter-node pipeline parallel (#40150) Signed-off-by: Fadi Arafeh --- .../device_communicators/cpu_communicator.py | 13 +++++++++++++ vllm/distributed/parallel_state.py | 6 ++++-- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/vllm/distributed/device_communicators/cpu_communicator.py b/vllm/distributed/device_communicators/cpu_communicator.py index 2bce5faa8b6..5dbf6b2057e 100644 --- a/vllm/distributed/device_communicators/cpu_communicator.py +++ b/vllm/distributed/device_communicators/cpu_communicator.py @@ -45,6 +45,9 @@ class CpuCommunicator(DeviceCommunicatorBase): unique_name, ) + # send/recv tensor_dict is only supported through the SHM communicator backend + self.supports_tensor_dict = isinstance(self.dist_module, _CPUSHMDistributed) + if self.use_all2all: if self.all2all_backend != "naive": # type: ignore[has-type] logger.warning( @@ -143,12 +146,22 @@ class CpuCommunicator(DeviceCommunicatorBase): tensor_dict: dict[str, torch.Tensor | Any], dst: int, ) -> None: + if not self.supports_tensor_dict: + raise NotImplementedError( + "CpuCommunicator does not support tensor dict fastpath with " + "torch.distributed backend." + ) return self.dist_module.send_tensor_dict(tensor_dict, dst) def recv_tensor_dict( self, src: int, ) -> dict[str, torch.Tensor | Any]: + if not self.supports_tensor_dict: + raise NotImplementedError( + "CpuCommunicator does not support tensor dict fastpath with " + "torch.distributed backend." + ) return self.dist_module.recv_tensor_dict(src) def dispatch_router_logits( diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index cef902d9e4e..473acb908b2 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -394,8 +394,10 @@ class GroupCoordinator: current_platform.is_tpu() or current_platform.use_custom_op_collectives() ) - self.use_cpu_custom_send_recv = current_platform.is_cpu() and hasattr( - torch.ops._C, "init_shm_manager" + self.use_cpu_custom_send_recv = ( + current_platform.is_cpu() + and self.device_communicator + and getattr(self.device_communicator, "supports_tensor_dict", False) ) def create_mq_broadcaster( From f774ba028afe08fab1ca96ecf7c0b40c0fa45357 Mon Sep 17 00:00:00 2001 From: Or Ozeri Date: Mon, 20 Apr 2026 12:53:51 +0300 Subject: [PATCH 137/150] [kv_offload+HMA][4/N]: Support sliding window lookup (#36645) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Or Ozeri Co-authored-by: Nicolò Lucchesi --- .../offloading_connector/test_scheduler.py | 160 +++++++++++++++++- .../unit/offloading_connector/utils.py | 19 ++- tests/v1/kv_offload/test_cpu_manager.py | 56 +++--- .../kv_connector/v1/offloading/scheduler.py | 52 +++++- vllm/v1/kv_offload/abstract.py | 20 +-- vllm/v1/kv_offload/cpu/manager.py | 15 +- vllm/v1/kv_offload/reuse_manager.py | 25 ++- 7 files changed, 269 insertions(+), 78 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 26bd01b138f..43d1fb94e70 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable +from unittest.mock import MagicMock import pytest @@ -10,8 +11,16 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import ( ) from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID from vllm.distributed.kv_events import BlockRemoved, BlockStored +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( + OffloadingConnectorScheduler, +) from vllm.v1.core.kv_cache_utils import BlockHash -from vllm.v1.kv_offload.abstract import OffloadingEvent +from vllm.v1.kv_offload.abstract import ( + OffloadingEvent, + OffloadingManager, + ReqContext, + get_offload_block_hash, +) from vllm.v1.request import RequestStatus @@ -105,8 +114,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): lambda keys, req_context: generate_store_output([]) ) runner.run(decoded_tokens=[EOS_TOKEN_ID]) - runner.manager.lookup.assert_called() - assert len(list(runner.manager.lookup.call_args.args[0])) == 1 + runner.manager.lookup.assert_called_once() # single block lookup with a hit runner.scheduler.reset_prefix_cache() @@ -114,7 +122,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.manager.prepare_store.side_effect = ( lambda keys, req_context: generate_store_output([]) ) - runner.manager.lookup.return_value = 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.run( decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2) ) @@ -126,7 +134,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.manager.prepare_store.side_effect = ( lambda keys, req_context: generate_store_output([]) ) - runner.manager.lookup.return_value = 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.run( decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5) ) @@ -210,7 +218,7 @@ def test_request_preemption(request_runner, async_scheduling: bool): # request should now return from preemption # re-load [0, ..., 8] from the CPU and store [9, 10, 11] - runner.manager.lookup.return_value = 3 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 3 runner.manager.prepare_store.side_effect = ( lambda keys, req_context: generate_store_output(keys) ) @@ -251,7 +259,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: # start a request to load the first block, but don't complete runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * offloaded_block_size) - runner.manager.lookup.return_value = 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.run( decoded_tokens=[], complete_transfers=False, @@ -263,7 +271,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: # start a new request to load the same first block runner.new_request(token_ids=[0] * offloaded_block_size) - runner.manager.lookup.return_value = 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.run( decoded_tokens=[], complete_transfers=False, @@ -311,7 +319,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool): # start a request to load the first block, but don't complete runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * offloaded_block_size) - runner.manager.lookup.return_value = 1 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.run( decoded_tokens=[], complete_transfers=False, @@ -336,3 +344,137 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool): # assert request is deleted assert req_id not in runner.scheduler.requests + + +# --------------------------------------------------------------------------- +# Unit tests for _maximal_prefix_lookup / _sliding_window_lookup +# --------------------------------------------------------------------------- + + +def _make_scheduler_with_lookup( + lookup_results: dict[int, bool | None], +) -> OffloadingConnectorScheduler: + """Create an OffloadingConnectorScheduler with a mocked manager.lookup.""" + manager = MagicMock(spec=OffloadingManager) + manager.lookup.side_effect = lambda key, req_context: lookup_results.get( + int(get_offload_block_hash(key).decode()), False + ) + + scheduler = object.__new__(OffloadingConnectorScheduler) + scheduler.manager = manager + return scheduler + + +_EMPTY_REQ_CTX = ReqContext() + + +class TestMaximalPrefixLookup: + def test_all_hit(self): + sched = _make_scheduler_with_lookup({1: True, 2: True}) + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2 + + def test_all_miss(self): + sched = _make_scheduler_with_lookup({}) + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 + + def test_partial_prefix(self): + sched = _make_scheduler_with_lookup({1: True, 2: True}) + assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2 + + def test_miss_then_hit(self): + sched = _make_scheduler_with_lookup({2: True}) + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 + + def test_single_hit(self): + sched = _make_scheduler_with_lookup({1: True}) + assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1 + + def test_empty(self): + sched = _make_scheduler_with_lookup({}) + assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0 + + def test_none_defers(self): + sched = _make_scheduler_with_lookup({1: None, 2: True}) + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None + + def test_none_after_hit_defers(self): + sched = _make_scheduler_with_lookup({1: True, 2: None}) + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None + + def test_none_stops_at_miss(self): + """None is treated as hit for iteration, but miss stops the scan.""" + sched = _make_scheduler_with_lookup({1: None, 2: False, 3: True}) + assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None + # lookup should have been called for blocks 1 and 2 (stops at miss) + assert sched.manager.lookup.call_count == 2 + + +class TestSlidingWindowLookup: + def test_all_hit_exact_window(self): + sched = _make_scheduler_with_lookup({1: True, 2: True}) + assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2 + + def test_all_miss(self): + sched = _make_scheduler_with_lookup({}) + assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 0 + + def test_window_at_end(self): + sched = _make_scheduler_with_lookup({2: True, 3: True}) + assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 2, _EMPTY_REQ_CTX) == 3 + + def test_window_in_middle(self): + sched = _make_scheduler_with_lookup({2: True, 3: True}) + assert ( + sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 3 + ) + + def test_no_full_window_falls_back_to_prefix(self): + sched = _make_scheduler_with_lookup({1: True, 2: True}) + assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 3, _EMPTY_REQ_CTX) == 2 + + def test_single_block_window(self): + sched = _make_scheduler_with_lookup({2: True, 3: True}) + assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 3 + + def test_gap_resets_consecutive(self): + sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True}) + # [1, 2, 3, 0, 4] — gap at 0 resets, window of 2 found at [2,3] + assert ( + sched._sliding_window_lookup(to_keys([1, 2, 3, 0, 4]), 2, _EMPTY_REQ_CTX) + == 3 + ) + + def test_window_prefers_rightmost(self): + sched = _make_scheduler_with_lookup({1: True, 2: True, 4: True, 5: True}) + # two valid windows: [1,2] at positions 0-1 and [4,5] at positions 3-4 + # scans right-to-left, finds [4,5] first + assert ( + sched._sliding_window_lookup(to_keys([1, 2, 3, 4, 5]), 2, _EMPTY_REQ_CTX) + == 5 + ) + + def test_prefix_fallback_with_gap(self): + sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True, 5: True}) + # window of 4 not found contiguously (gap at 1) + assert ( + sched._sliding_window_lookup(to_keys([2, 1, 3, 4, 5]), 4, _EMPTY_REQ_CTX) + == 1 + ) + + def test_empty(self): + sched = _make_scheduler_with_lookup({}) + assert sched._sliding_window_lookup([], 1, _EMPTY_REQ_CTX) == 0 + + def test_none_defers(self): + sched = _make_scheduler_with_lookup({1: True, 2: None}) + assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) is None + + def test_none_with_full_window_still_defers(self): + """Even if a real window is found after a None, result is deferred.""" + # Scan right-to-left: 4(True), 3(None) resets, 2(True), 1(True) = window + # but block 3 was None so defer_lookup is set + sched = _make_scheduler_with_lookup({1: True, 2: True, 3: None, 4: True}) + assert ( + sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) + is None + ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index aaf9152a43c..0888c061536 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -56,8 +56,12 @@ from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager -def to_keys(int_ids: list[int]) -> list[OffloadKey]: - return [make_offload_key(str(i).encode(), 0) for i in int_ids] +def to_key(int_hash: int) -> OffloadKey: + return make_offload_key(str(int_hash).encode(), 0) + + +def to_keys(int_hashes: list[int]) -> list[OffloadKey]: + return [to_key(i) for i in int_hashes] class MockLoadStoreSpec(LoadStoreSpec): @@ -116,6 +120,7 @@ class MockOffloadingSpec(OffloadingSpec): self.manager = MagicMock(spec=OffloadingManager) self.manager.lookup.return_value = 0 self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys) + self.manager.lookup.return_value = False self.handler = MockOffloadingHandler() def get_manager(self) -> OffloadingManager: @@ -228,14 +233,14 @@ class RequestRunner: self.scheduler_connector: OffloadingConnector = scheduler_connector # extract mocked OffloadingManager of scheduler connector - connector_scheduler = scheduler_connector.connector_scheduler - assert connector_scheduler is not None - manager = connector_scheduler.manager + self.connector_scheduler = scheduler_connector.connector_scheduler + assert self.connector_scheduler is not None + manager = self.connector_scheduler.manager assert isinstance(manager, MagicMock) self.manager: MagicMock = manager - assert len(connector_scheduler.config.kv_group_configs) == 1 - kv_group_config = connector_scheduler.config.kv_group_configs[0] + assert len(self.connector_scheduler.config.kv_group_configs) == 1 + kv_group_config = self.connector_scheduler.config.kv_group_configs[0] assert kv_group_config.gpu_block_size == gpu_block_size assert kv_group_config.offloaded_block_size == offloaded_block_size diff --git a/tests/v1/kv_offload/test_cpu_manager.py b/tests/v1/kv_offload/test_cpu_manager.py index 651ea091ae6..733f9bf519e 100644 --- a/tests/v1/kv_offload/test_cpu_manager.py +++ b/tests/v1/kv_offload/test_cpu_manager.py @@ -35,8 +35,12 @@ class ExpectedPrepareStoreOutput: evicted_keys: list[int] -def to_keys(int_ids: list[int]) -> list[OffloadKey]: - return [make_offload_key(str(i).encode(), 0) for i in int_ids] +def to_key(int_hash: int) -> OffloadKey: + return make_offload_key(str(int_hash).encode(), 0) + + +def to_keys(int_hashes: list[int]) -> list[OffloadKey]: + return [to_key(i) for i in int_hashes] def verify_store_output( @@ -136,7 +140,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): manager.complete_store(to_keys([2, 3, 4, 5])) # block 2 must still be present in the cache - assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1 + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True def test_cpu_manager(): @@ -160,7 +164,8 @@ def test_cpu_manager(): ) # lookup [1, 2] -> not ready - assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False # no events so far assert list(cpu_manager.take_events()) == [] @@ -170,9 +175,9 @@ def test_cpu_manager(): verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] - assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1 - assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2 - assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False # prepare store [2, 3, 4, 5] -> evicts [1] prepare_store_output = cpu_manager.prepare_store( @@ -196,6 +201,14 @@ def test_cpu_manager(): # complete store [2, 3, 4, 5] cpu_manager.complete_store(to_keys([2, 3, 4, 5])) + # lookup (now that we have [2, 3, 4, 5]) + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(4), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is False + # prepare load [2, 3] prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX) verify_load_output(prepare_load_output, [1, 2]) @@ -238,8 +251,8 @@ def test_cpu_manager(): cpu_manager.complete_store(to_keys([7, 9]), success=False) # assert [7] is still stored, but [9] is not - assert cpu_manager.lookup(to_keys([7]), _EMPTY_REQ_CTX) == 1 - assert cpu_manager.lookup(to_keys([9]), _EMPTY_REQ_CTX) == 0 + assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(9), _EMPTY_REQ_CTX) is False verify_events( cpu_manager.take_events(), @@ -284,7 +297,8 @@ class TestARCPolicy: ) # lookup [1, 2] -> not ready - assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False # no events so far assert list(cpu_manager.take_events()) == [] @@ -294,9 +308,9 @@ class TestARCPolicy: verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] - assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1 - assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2 - assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False # blocks should be in T1 (recent) assert len(arc_policy.t1) == 2 @@ -500,7 +514,7 @@ class TestARCPolicy: cpu_manager.complete_store(to_keys([5]), success=False) # block 5 should not be in cache - assert cpu_manager.lookup(to_keys([5]), _EMPTY_REQ_CTX) == 0 + assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is False # block 5 should not be in T1 or T2 assert to_keys([5])[0] not in arc_policy.t1 assert to_keys([5])[0] not in arc_policy.t2 @@ -541,8 +555,8 @@ class TestARCPolicy: cpu_manager.complete_store(to_keys([6])) # verify blocks 2, 3 (in T2) are still present - assert cpu_manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1 - assert cpu_manager.lookup(to_keys([3]), _EMPTY_REQ_CTX) == 1 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True # verify events events = list(cpu_manager.take_events()) @@ -562,7 +576,8 @@ def test_filter_reused_manager(): ) # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet - assert manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 + assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False # prepare store [1, 2] -> should be filtered prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) @@ -570,7 +585,7 @@ def test_filter_reused_manager(): assert prepare_store_output.keys_to_store == [] # Lookup [1] -> 2nd time, eligible now - assert manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 0 + assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) @@ -579,12 +594,13 @@ def test_filter_reused_manager(): # Lookup [3, 4] -> 1st time # (evicts [2] from tracker since max_size is 3 and tracker has [1]) - assert manager.lookup(to_keys([3, 4]), _EMPTY_REQ_CTX) == 0 + assert manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False + assert manager.lookup(to_key(4), _EMPTY_REQ_CTX) is False # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4]) assert to_keys([2])[0] not in manager.counts # Lookup [2] again -> (this adds [2] back to the tracker as 1st time) - assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 0 + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False # Verify [2] was re-added with count=1 (not eligible yet) assert manager.counts.get(to_keys([2])[0]) == 1 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index c5272ea2778..cd5a4f113dc 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import defaultdict -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from dataclasses import dataclass, field from itertools import islice from typing import Any, NamedTuple @@ -132,6 +132,49 @@ class OffloadingConnectorScheduler: self._reqs_being_stored = defaultdict[ReqId, set[OffloadKey]](set) self._reqs_being_loaded = defaultdict[ReqId, set[OffloadKey]](set) + def _maximal_prefix_lookup( + self, keys: Iterable[OffloadKey], req_context: ReqContext + ) -> int | None: + """Find the length of the maximal prefix of offloaded blocks.""" + hit_count = 0 + defer_lookup = False + for key in keys: + result = self.manager.lookup(key, req_context) + if result is None: + defer_lookup = True + # continue lookup to allow manager to kick-off async lookups + # for all blocks (until a miss is detected) + result = True + if not result: + break + hit_count += 1 + return hit_count if not defer_lookup else None + + def _sliding_window_lookup( + self, + keys: Sequence[OffloadKey], + sliding_window_size: int, + req_context: ReqContext, + ) -> int | None: + """Find the maximal ending position of consecutive offloaded blocks + within a sliding window.""" + defer_lookup = False + consecutive_hits = 0 + for idx in range(len(keys) - 1, -1, -1): + result = self.manager.lookup(keys[idx], req_context) + if result is None: + defer_lookup = True + # continue lookup to allow manager to kick-off async lookups + # for all blocks (until a hit is detected) + result = False + if not result: + consecutive_hits = 0 + else: + consecutive_hits += 1 + if consecutive_hits == sliding_window_size: + return idx + sliding_window_size if not defer_lookup else None + return consecutive_hits if not defer_lookup else None + def get_num_new_matched_tokens( self, request: Request, num_computed_tokens: int ) -> tuple[int | None, bool]: @@ -184,9 +227,10 @@ class OffloadingConnectorScheduler: return 0, False start_block_idx = num_computed_tokens // group_config.offloaded_block_size - hits = self.manager.lookup( - offload_keys[start_block_idx:], - req_status.req_context, + # Full attention relays on all previous KV cache blocks. + # Thus, we search for a maximal prefix of KV cache which are all cached. + hits = self._maximal_prefix_lookup( + offload_keys[start_block_idx:], req_status.req_context ) if hits is None: # indicates a lookup that should be tried later diff --git a/vllm/v1/kv_offload/abstract.py b/vllm/v1/kv_offload/abstract.py index 39ee360e8f6..8f809ceaa08 100644 --- a/vllm/v1/kv_offload/abstract.py +++ b/vllm/v1/kv_offload/abstract.py @@ -7,8 +7,7 @@ This class runs in the scheduler, tracks which blocks are offloaded and their address. The class provides the following primitives: - lookup() - find the length of the maximal series of blocks, - starting from the first one, that are all offloaded. + lookup() - check whether a single block is offloaded and ready. prepare_load() - prepare given blocks to be read. The given blocks will be protected from eviction. This function returns a LoadSpec which encapsulates @@ -91,23 +90,18 @@ class OffloadingEvent: class OffloadingManager(ABC): @abstractmethod - def lookup( - self, - keys: Iterable[OffloadKey], - req_context: ReqContext, - ) -> int | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: """ - Finds the length of the maximal series of blocks, starting from the - first one, that are all offloaded. + Checks whether a single block is offloaded and ready to be read. Args: - keys: the keys identifying the blocks to lookup. + key: the key identifying the block to lookup. req_context: per-request context (e.g. kv_transfer_params). Returns: - An integer representing the maximal number of blocks that - are currently offloaded, or None if the lookup should be retried - later. Returning None will delay the request handling by the vLLM + True if the block is offloaded and ready, False if not, + or None if the lookup should be retried later. + Returning None will delay the request handling by the vLLM scheduler. """ pass diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 5ae7454430f..fcfaa919a3b 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -84,18 +84,9 @@ class CPUOffloadingManager(OffloadingManager): # --- OffloadingManager interface --- - def lookup( - self, - keys: Iterable[OffloadKey], - req_context: ReqContext, - ) -> int | None: - hit_count = 0 - for key in keys: - block = self._policy.get(key) - if block is None or not block.is_ready: - break - hit_count += 1 - return hit_count + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + block = self._policy.get(key) + return block is not None and block.is_ready def prepare_load( self, diff --git a/vllm/v1/kv_offload/reuse_manager.py b/vllm/v1/kv_offload/reuse_manager.py index a9650e38c51..96b8f969e75 100644 --- a/vllm/v1/kv_offload/reuse_manager.py +++ b/vllm/v1/kv_offload/reuse_manager.py @@ -27,8 +27,9 @@ class FilterReusedOffloadingManager(OffloadingManager): All methods are delegated to the *backing* manager. Two methods are intercepted: - * ``lookup`` — records each visited key in an internal LRU counter. * ``prepare_store`` — filters out keys that have not yet + * ``lookup`` — records the visited key in an internal LRU + counter, then delegates to the backing manager. crossed the threshold *before* calling the backing ``prepare_store``. @@ -66,18 +67,16 @@ class FilterReusedOffloadingManager(OffloadingManager): # Intercepted methods # ------------------------------------------------------------------ - def lookup(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> int | None: - """Record each key, then delegate lookup to backing manager.""" - keys = list(keys) - for key in keys: - if key in self.counts: - self.counts.move_to_end(key) - self.counts[key] += 1 - else: - if len(self.counts) >= self.max_tracker_size: - self.counts.popitem(last=False) # evict LRU - self.counts[key] = 1 - return self._backing.lookup(keys, req_context) + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + """Record the key, then delegate lookup to backing manager.""" + if key in self.counts: + self.counts.move_to_end(key) + self.counts[key] += 1 + else: + if len(self.counts) >= self.max_tracker_size: + self.counts.popitem(last=False) # evict LRU + self.counts[key] = 1 + return self._backing.lookup(key, req_context) def prepare_store( self, keys: Iterable[OffloadKey], req_context: ReqContext From 50dd4cb42726777635acda9ed4f5440ae4a2e281 Mon Sep 17 00:00:00 2001 From: Ilya Markov Date: Mon, 20 Apr 2026 12:24:23 +0200 Subject: [PATCH 138/150] [EPLB] Add nixl-based eplb communicator (#36276) Signed-off-by: ilmarkov Signed-off-by: Markov Ilya --- docs/serving/expert_parallel_deployment.md | 1 + tests/distributed/test_eplb_execute.py | 17 +- vllm/config/parallel.py | 3 +- vllm/distributed/eplb/eplb_communicator.py | 446 +++++++++++++++++- vllm/distributed/eplb/rebalance_execute.py | 34 +- .../kv_transfer/kv_connector/v1/nixl/stats.py | 4 +- .../kv_transfer/kv_connector/v1/nixl/utils.py | 46 -- .../kv_connector/v1/nixl/worker.py | 3 +- vllm/distributed/nixl_utils.py | 54 +++ vllm/v1/worker/gpu_model_runner.py | 28 +- 10 files changed, 556 insertions(+), 80 deletions(-) create mode 100644 vllm/distributed/nixl_utils.py diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index d75ae7feb49..fef4df770fa 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -153,6 +153,7 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` | | `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` | | `policy` | The policy type for expert parallel load balancing | `"default"` | +| `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`, or `null` (auto) | `null` | For example: diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index 7256ad56a3f..8a46087e991 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -9,7 +9,10 @@ import torch import torch.distributed from vllm.config import VllmConfig, set_current_vllm_config -from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator +from vllm.distributed.eplb.eplb_communicator import ( + create_eplb_communicator, + has_nixl, +) from vllm.distributed.eplb.rebalance_execute import ( move_from_buffer, rearrange_expert_weights_inplace, @@ -527,7 +530,9 @@ def _test_rearrange_expert_weights_with_redundancy( (4, 8, 8, 16), ], ) -@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"]) +@pytest.mark.parametrize( + "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"] +) def test_rearrange_expert_weights_with_redundancy( world_size, num_layers, @@ -537,6 +542,8 @@ def test_rearrange_expert_weights_with_redundancy( ): """Test the functionality of rearranging expert weights with redundancy.""" + if eplb_communicator == "nixl" and not has_nixl(): + pytest.skip("NIXL is not available") if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") distributed_run( @@ -633,7 +640,9 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: (2, 2, 2, 3), ], ) -@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"]) +@pytest.mark.parametrize( + "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"] +) def test_async_transfer_layer_without_mtp( world_size: int, num_layers: int, @@ -643,6 +652,8 @@ def test_async_transfer_layer_without_mtp( ): """Exercise async EPLB transfer path without MTP/spec decode.""" + if eplb_communicator == "nixl" and not has_nixl(): + pytest.skip("NIXL is not available") if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index a42b8422ef3..afd0d1dd501 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -36,7 +36,7 @@ DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"] DataParallelBackend = Literal["ray", "mp"] EPLBPolicyOption = Literal["default"] DCPCommBackend = Literal["ag_rs", "a2a"] -EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "pynccl"] +EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "nixl", "pynccl"] All2AllBackend = Literal[ "naive", "pplx", @@ -90,6 +90,7 @@ class EPLBConfig: Backend for EPLB expert weight communication: - "torch_nccl": Use torch.distributed on the device process group - "torch_gloo": Use torch.distributed gloo with CPU staging + - "nixl": Use NIXL/ RIXL with staged send/recv buffers - "pynccl": Use PyNccl send/recv - None: Auto-select backend ("torch_gloo" for async, "torch_nccl" for sync) """ diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 059d6b7cffb..6ff41272fce 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -4,8 +4,12 @@ EPLB communicator implementations and factory. """ +import contextlib +import time +import uuid from abc import ABC, abstractmethod from collections.abc import Sequence +from datetime import timedelta import torch from torch.distributed import ( @@ -18,13 +22,27 @@ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.distributed.device_communicators.pynccl_wrapper import ( ncclDataTypeEnum, ) -from vllm.distributed.parallel_state import GroupCoordinator, is_local_first_rank +from vllm.distributed.nixl_utils import ( + NixlWrapper, + nixl_agent_config, +) +from vllm.distributed.parallel_state import ( + GroupCoordinator, + get_pp_group, + is_local_first_rank, +) from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator from vllm.logger import init_logger +from vllm.platforms import current_platform logger = init_logger(__name__) +def has_nixl() -> bool: + """Whether the optional NIXL / RIXL package is available.""" + return NixlWrapper is not None + + class EplbCommunicator(ABC): """Abstract EPLB communicator for expert weight transfers.""" @@ -40,6 +58,12 @@ class EplbCommunicator(ABC): def execute(self) -> None: pass + @property + def needs_profile_buffer_reservation(self) -> bool: + """Whether the profile path must run a dummy collective operation to reserve + communication buffers.""" + return True + def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None: self._cuda_stream = cuda_stream @@ -167,6 +191,385 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator): dst_tensor.copy_(cpu_tensor, non_blocking=True) +class NixlEplbCommunicator(EplbCommunicator): + """EPLB communicator backed by NIXL READ transfers.""" + + def __init__( + self, + cpu_group: ProcessGroup, + expert_weights: Sequence[torch.Tensor], + cuda_stream: torch.cuda.Stream | None = None, + ) -> None: + assert expert_weights, "NixlEplbCommunicator requires non-empty expert_weights." + if NixlWrapper is None: + raise RuntimeError("NIXL/ RIXL is unavailable.") + self._cpu_group = cpu_group + self._cuda_stream = cuda_stream + self._world_size = cpu_group.size() + self._rank = cpu_group.rank() + self._send_tensors: dict[torch.dtype, list[list[torch.Tensor]]] = {} + self._recv_tensors: dict[torch.dtype, list[list[torch.Tensor]]] = {} + self._dtypes: list[torch.dtype] = [] + self._device = expert_weights[0].device + for tensor in expert_weights: + assert tensor.device == self._device, ( + "All local EPLB tensors are expected to be on the same device: " + f"expected={self._device}, got={tensor.device}" + ) + if tensor.dtype not in self._dtypes: + self._dtypes.append(tensor.dtype) + + config = ( + nixl_agent_config(capture_telemetry=False) + if nixl_agent_config is not None + else None + ) + self._nixl_wrapper = NixlWrapper(self._make_agent_name(), config) + self._nixl_memory_type = "VRAM" + self._registered_desc: object | None = None + self._remote_agents: dict[int, str] = {} + self._remote_send_meta: dict[int, tuple[int, int, int]] = {} + self._send_buffer: torch.Tensor = torch.empty(0) + self._recv_buffer: torch.Tensor = torch.empty(0) + self._peer_partition_bytes: int = 0 + self._dtype_max_bytes: dict[torch.dtype, int] = {} + self._cuda_device_id = int(self._device.index or 0) + self._xfer_cache: dict[tuple[int, int, int], tuple[int, int, int]] = {} + self._init_step("buffers", self._init_registered_buffers, expert_weights) + self._init_step("agents", self._init_remote_agents) + self._init_step("send meta", self._exchange_remote_send_meta) + self._log_initialized() + + @property + def needs_profile_buffer_reservation(self) -> bool: + return False + + @staticmethod + def _init_step(name: str, fn: object, *args: object, **kwargs: object) -> None: + try: + fn(*args, **kwargs) # type: ignore[operator] + except Exception as exc: + raise RuntimeError(f"NIXL EPLB init failed: {name}") from exc + + def _make_agent_name(self) -> str: + """Build a deployment-unique nixl agent name.""" + pp_size = get_pp_group().world_size + pp_suffix = f"-pp{get_pp_group().rank_in_group}" if pp_size > 1 else "" + uid = uuid.uuid4().hex[:8] + return f"eplb-{self._rank}{pp_suffix}-{uid}" + + def _get_peer_buckets( + self, + bucket_map: dict[torch.dtype, list[list[torch.Tensor]]], + dtype: torch.dtype, + ) -> list[list[torch.Tensor]]: + peer_buckets = bucket_map.get(dtype) + if peer_buckets is None: + peer_buckets = [[] for _ in range(self._world_size)] + bucket_map[dtype] = peer_buckets + return peer_buckets + + def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None: + assert dst_rank != self._rank, ( + "EPLB communicator should not enqueue same-rank sends: " + f"rank={self._rank}, dst_rank={dst_rank}" + ) + self._get_peer_buckets(self._send_tensors, tensor.dtype)[dst_rank].append( + tensor + ) + + def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None: + assert src_rank != self._rank, ( + "EPLB communicator should not enqueue same-rank recvs: " + f"rank={self._rank}, src_rank={src_rank}" + ) + self._get_peer_buckets(self._recv_tensors, tensor.dtype)[src_rank].append( + tensor + ) + + def _init_remote_agents(self) -> None: + local_metadata = self._nixl_wrapper.get_agent_metadata() + gathered_metadata: list[bytes | None] = [None] * self._world_size + torch.distributed.all_gather_object( + gathered_metadata, local_metadata, group=self._cpu_group + ) + for peer in range(self._world_size): + if peer == self._rank: + continue + peer_metadata = gathered_metadata[peer] + assert peer_metadata is not None + self._remote_agents[peer] = self._nixl_wrapper.add_remote_agent( + peer_metadata + ) + + def _init_registered_buffers(self, expert_weights: Sequence[torch.Tensor]) -> None: + total_max_bytes = 0 + for dtype in self._dtypes: + max_numel = max( + sum(t.numel() for t in expert_weights if t.dtype == dtype), 1 + ) + max_bytes = max_numel * dtype.itemsize + self._dtype_max_bytes[dtype] = max_bytes + total_max_bytes += max_bytes + + self._peer_partition_bytes = total_max_bytes + + # The send buffer needs world_size partitions because remote peers + # READ from fixed offsets (rank * partition_bytes). + # This allocates world_size * partition_bytes + # which can cause OOM on large models. + # TODO(ilmarkov): shrink to const * partition_bytes and execute + # communication in multiple steps dealing with the worst case. + send_total_bytes = self._peer_partition_bytes * self._world_size + + self._send_buffer = torch.empty( + send_total_bytes, device=self._device, dtype=torch.uint8 + ) + self._recv_buffer = torch.empty( + self._peer_partition_bytes, device=self._device, dtype=torch.uint8 + ) + + descs = self._nixl_wrapper.get_reg_descs([self._send_buffer, self._recv_buffer]) + self._nixl_wrapper.register_memory(descs) + self._registered_desc = descs + + def _exchange_remote_send_meta(self) -> None: + """Exchange send-buffer metadata so each rank can build dynamic + descriptors at execute time.""" + local_meta: tuple[int, int, int] = ( + self._send_buffer.data_ptr(), + self._peer_partition_bytes, + self._cuda_device_id, + ) + gathered_meta: list[tuple[int, int, int] | None] = [None] * self._world_size + torch.distributed.all_gather_object( + gathered_meta, local_meta, group=self._cpu_group + ) + + for peer in self._remote_agents: + peer_meta = gathered_meta[peer] + assert peer_meta is not None + self._remote_send_meta[peer] = peer_meta + + @staticmethod + def _pack_send_buffer( + peer_tensors: list[torch.Tensor], + send_buffer: torch.Tensor, + byte_offset: int, + ) -> int: + """ + Returns the byte offset after the last written byte. + """ + for tensor in peer_tensors: + raw = tensor.reshape(-1).view(torch.uint8) + if raw.numel() == 0: + continue + send_buffer[byte_offset : byte_offset + raw.numel()].copy_( + raw, non_blocking=True + ) + byte_offset += raw.numel() + return byte_offset + + @staticmethod + def _unpack_recv_buffer( + recv_buffer: torch.Tensor, + peer_tensors: list[torch.Tensor], + byte_offset: int, + ) -> int: + """ + Returns the byte offset after the last read byte. + """ + for tensor in peer_tensors: + num_bytes = tensor.numel() * tensor.element_size() + if num_bytes == 0: + continue + tensor.reshape(-1).view(torch.uint8).copy_( + recv_buffer[byte_offset : byte_offset + num_bytes], + non_blocking=True, + ) + byte_offset += num_bytes + return byte_offset + + def _release_all_cached_handles(self) -> None: + """Best-effort release of every cached dlist and xfer handle.""" + for local_dlist, remote_dlist, xfer in self._xfer_cache.values(): + for release_fn, handle in ( + (self._nixl_wrapper.release_xfer_handle, xfer), + (self._nixl_wrapper.release_dlist_handle, local_dlist), + (self._nixl_wrapper.release_dlist_handle, remote_dlist), + ): + with contextlib.suppress(Exception): + release_fn(handle) + self._xfer_cache.clear() + + def _wait_for_all_transfers(self, handles: list[int]) -> None: + pending = set(handles) + while pending: + completed: list[int] = [] + for handle in pending: + state = self._nixl_wrapper.check_xfer_state(handle) + if state == "DONE": + completed.append(handle) + continue + if state != "PROC": + raise RuntimeError(f"NIXL transfer failed with state={state}") + for handle in completed: + pending.remove(handle) + if pending: + time.sleep(0.0005) + + def _get_or_create_xfer(self, src: int, total_bytes: int, recv_offset: int) -> int: + """Return a cached xfer handle or create and cache a new one.""" + key = (src, total_bytes, recv_offset) + cached = self._xfer_cache.get(key) + if cached is not None: + return cached[2] + + recv_base = self._recv_buffer.data_ptr() + local_desc = self._nixl_wrapper.get_xfer_descs( + [ + ( + recv_base + recv_offset, + total_bytes, + self._cuda_device_id, + ) + ], + self._nixl_memory_type, + ) + local_handle = self._nixl_wrapper.prep_xfer_dlist( + "NIXL_INIT_AGENT", + local_desc, + ) + + remote_base, remote_part_bytes, remote_dev = self._remote_send_meta[src] + agent_name = self._remote_agents[src] + remote_desc = self._nixl_wrapper.get_xfer_descs( + [ + ( + remote_base + self._rank * remote_part_bytes, + total_bytes, + remote_dev, + ) + ], + self._nixl_memory_type, + ) + remote_handle = self._nixl_wrapper.prep_xfer_dlist( + agent_name, + remote_desc, + ) + + xfer_handle = self._nixl_wrapper.make_prepped_xfer( + "READ", + local_handle, + [0], + remote_handle, + [0], + ) + self._xfer_cache[key] = (local_handle, remote_handle, xfer_handle) + return xfer_handle + + def execute(self) -> None: + xfer_handles: list[int] = [] + try: + # Phase 1: pack send buffers. + with torch.cuda.stream(self._cuda_stream): + for dst in range(self._world_size): + byte_offset = dst * self._peer_partition_bytes + for dtype in self._dtypes: + peer_tensors = self._send_tensors.get( + dtype, [[] for _ in range(self._world_size)] + )[dst] + actual_bytes = sum( + t.numel() * t.element_size() for t in peer_tensors + ) + if actual_bytes > self._dtype_max_bytes[dtype]: + raise RuntimeError( + "NIXL EPLB send overflow for dtype " + f"{dtype}: peer={dst}, " + f"required={actual_bytes}, " + f"capacity={self._dtype_max_bytes[dtype]}" + ) + byte_offset = self._pack_send_buffer( + peer_tensors, + self._send_buffer, + byte_offset, + ) + + # Ensure all packed data is visible in device memory before pulls. + if self._cuda_stream is not None: + self._cuda_stream.synchronize() + else: + torch.cuda.current_stream().synchronize() + # READ is receiver-initiated; synchronize all ranks before transfer. + # We use monitored_barrier so a rank that crashes or exits early + # produces a diagnostic timeout instead of a silent hang. + torch.distributed.monitored_barrier( + group=self._cpu_group, + timeout=timedelta(minutes=5), + ) + + # Phase 2: look up or create descriptors and issue all READs. + # Data from all peers is packed sequentially into the single + # partition-sized recv buffer at running offsets. + recv_offsets: dict[int, int] = {} + recv_offset = 0 + for src in range(self._world_size): + if src == self._rank: + continue + actual_total_bytes = 0 + for dtype in self._dtypes: + peer_tensors = self._recv_tensors.get( + dtype, [[] for _ in range(self._world_size)] + )[src] + actual_total_bytes += sum( + t.numel() * t.element_size() for t in peer_tensors + ) + if actual_total_bytes == 0: + continue + + recv_offsets[src] = recv_offset + xfer_handle = self._get_or_create_xfer( + src, actual_total_bytes, recv_offset + ) + self._nixl_wrapper.transfer(xfer_handle) + xfer_handles.append(xfer_handle) + recv_offset += actual_total_bytes + + # Phase 3: single wait for all in-flight transfers, then unpack. + self._wait_for_all_transfers(xfer_handles) + + with torch.cuda.stream(self._cuda_stream): + for src, offset in recv_offsets.items(): + byte_offset = offset + for dtype in self._dtypes: + peer_tensors = self._recv_tensors.get( + dtype, [[] for _ in range(self._world_size)] + )[src] + byte_offset = self._unpack_recv_buffer( + self._recv_buffer, + peer_tensors, + byte_offset, + ) + except Exception: + self._release_all_cached_handles() + raise + finally: + self._send_tensors.clear() + self._recv_tensors.clear() + + def __del__(self) -> None: + try: + self._release_all_cached_handles() + if self._registered_desc is not None: + self._nixl_wrapper.deregister_memory(self._registered_desc) + self._registered_desc = None + for agent_name in self._remote_agents.values(): + self._nixl_wrapper.remove_remote_agent(agent_name) + self._remote_agents.clear() + except Exception as e: + logger.warning("Error during NixlEplbCommunicator cleanup: %s", e) + + class PyNcclEplbCommunicator(EplbCommunicator): """EPLB communicator backed by PyNcclCommunicator using ncclSend/ncclRecv.""" @@ -204,6 +607,24 @@ def create_eplb_communicator( backend: str | None, expert_weights: Sequence[torch.Tensor], ) -> EplbCommunicator: + """Create an EPLB communicator for the given backend. + + Args: + group_coordinator: Process-group coordinator that provides the + device and CPU communication groups. + backend: Communicator backend name (``"torch_nccl"``, + ``"torch_gloo"``, ``"pynccl"``, or ``"nixl"``). + Falls back to ``"torch_nccl"`` when *None*. + Stateless (elastic EP) groups only support ``"torch_nccl"`` + and ``"pynccl"``; ``"torch_nccl"`` is silently promoted to + ``"pynccl"`` in that case. When tensors reside on CPU, + ``"torch_gloo"`` or ``"torch_nccl"`` are used via the CPU + process group. + expert_weights: Expert weight tensors from *one* MoE layer. + NixlEplbCommunicator pre-allocates send/recv buffers sized + to this layer, so all other MoE layers must have the same + tensor count, shapes, and dtypes. + """ # Keep a safe default for callers that have not resolved communicator yet. if backend is None: backend = "torch_nccl" @@ -256,7 +677,7 @@ def create_eplb_communicator( if backend not in ("torch_nccl", "pynccl"): raise ValueError( f"Elastic EP requires 'torch_nccl' or 'pynccl' EPLB communicator " - f"(got '{backend}'). torch_gloo is not supported with stateless groups." + f"(got '{backend}')." ) if backend == "torch_nccl": logger.warning( @@ -266,7 +687,26 @@ def create_eplb_communicator( backend = "pynccl" return _create_pynccl() - if backend == "torch_gloo": + if backend == "nixl": + if not has_nixl(): + raise RuntimeError( + "EPLB communicator 'nixl' requested but NIXL is unavailable." + ) + if not (current_platform.is_cuda_alike() and tensor_device_type != "cpu"): + raise RuntimeError( + "EPLB communicator 'nixl' supports only cuda-like devices " + f"(got {tensor_device_type})." + ) + try: + return NixlEplbCommunicator( + cpu_group=group_coordinator.cpu_group, + expert_weights=expert_weights, + ) + except Exception as exc: + raise RuntimeError( + f"Failed to initialize NixlEplbCommunicator ({exc})." + ) from exc + elif backend == "torch_gloo": return TorchDistGlooStagedEplbCommunicator( cpu_group=group_coordinator.cpu_group, ) diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index 6670ccc6ee6..6c8a2dcc319 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -541,23 +541,29 @@ def rearrange_expert_weights_inplace( assert num_physical_experts == ep_size * num_local_physical_experts first_layer_weights = list(expert_weights[0]) + + if is_profile: + if communicator.needs_profile_buffer_reservation: + # Reserve NCCL communication buffers via a dummy all_gather. + # Backends that pre-allocate their own transfer buffers + # skip this to avoid the extra memory spike during profiling. + weights_buffer: list[torch.Tensor] = [ + torch.empty_like(w) for w in first_layer_weights + ] + for weight, buffer in zip(expert_weights[0], weights_buffer): + dummy_recv_buffer = [buffer for _ in range(ep_size)] + torch.distributed.barrier() + all_gather( + dummy_recv_buffer, + weight, + group=ep_group, + ) + return + # Buffers to hold the expert weights during the exchange. # NOTE: Currently we assume the same weights across different layers # have the same shape. - weights_buffer: list[torch.Tensor] = [ - torch.empty_like(w) for w in first_layer_weights - ] - if is_profile: - # Reserve communication buffers via a minimal dummy all_gather on first layer - for weight, buffer in zip(expert_weights[0], weights_buffer): - dummy_recv_buffer = [buffer for _ in range(ep_size)] - torch.distributed.barrier() - all_gather( - dummy_recv_buffer, - weight, - group=ep_group, - ) - return + weights_buffer = [torch.empty_like(w) for w in first_layer_weights] # NOTE(bowen): We need this synchronize to run, but I don't know why. # If you figure out the reason, please let me know -- thank you! diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py index fde99c1e2c2..65c553cfec3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py @@ -15,9 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( PromMetric, PromMetricT, ) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( - nixlXferTelemetry, -) +from vllm.distributed.nixl_utils import nixlXferTelemetry from vllm.v1.metrics.utils import create_metric_per_engine diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py index 514214347ae..d0b72464a27 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py @@ -3,60 +3,14 @@ """Shared constants, lazy imports and helpers for the NIXL connector.""" import contextlib -import os -import sys from collections.abc import Iterator from typing import Any import zmq -from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.network_utils import make_zmq_socket -logger = init_logger(__name__) - - -# Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used -try: - if "UCX_RCACHE_MAX_UNRELEASED" not in os.environ: - # avoid a memory leak in UCX when using NIXL on some models - # see: https://github.com/vllm-project/vllm/issues/24264 - if "nixl" in sys.modules or "rixl" in sys.modules: - logger.warning( - "NIXL was already imported, we can't reset UCX_RCACHE_MAX_UNRELEASED. " - "Please set it to '1024' manually." - ) - else: - logger.info( - "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare " - "memory leak in UCX when using NIXL." - ) - os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024" - - if not current_platform.is_rocm(): - from nixl._api import nixl_agent as NixlWrapper - from nixl._bindings import nixlXferTelemetry - else: - from rixl._api import nixl_agent as NixlWrapper - from rixl._bindings import nixlXferTelemetry - - logger.info("NIXL is available") -except ImportError: - logger.warning("NIXL is not available") - NixlWrapper = None - nixlXferTelemetry = None - - -try: - if not current_platform.is_rocm(): - from nixl._api import nixl_agent_config - else: - from rixl._api import nixl_agent_config -except ImportError: - nixl_agent_config = None - logger.warning("NIXL agent config is not available") - # Supported platforms and types of kv transfer buffer. # {device: tuple of supported kv buffer types} _NIXL_SUPPORTED_DEVICE = { diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py index 45aa33033e7..8157df63d1b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py @@ -45,8 +45,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( _NIXL_SUPPORTED_DEVICE, - NixlWrapper, - nixl_agent_config, zmq_ctx, ) from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( @@ -54,6 +52,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import compute_mamba_phys_ratio, derive_mamba_conv_split, ) +from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config from vllm.distributed.parallel_state import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, diff --git a/vllm/distributed/nixl_utils.py b/vllm/distributed/nixl_utils.py new file mode 100644 index 00000000000..b2b433339be --- /dev/null +++ b/vllm/distributed/nixl_utils.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import sys + +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +if "UCX_RCACHE_MAX_UNRELEASED" not in os.environ: + if "nixl" in sys.modules or "rixl" in sys.modules: + logger.warning_once( + "NIXL was already imported, we can't reset " + "UCX_RCACHE_MAX_UNRELEASED. " + "Please set it to '1024' manually." + ) + else: + logger.info_once( + "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare " + "memory leak in UCX when using NIXL." + ) + os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024" + +try: + if current_platform.is_cuda(): + from nixl._api import nixl_agent as NixlWrapper + else: + from rixl._api import nixl_agent as NixlWrapper + + logger.info_once("NIXL is available") +except ImportError: + logger.warning_once("NIXL is not available") + NixlWrapper = None # type: ignore[assignment, misc] + +try: + if current_platform.is_cuda(): + from nixl._api import nixl_agent_config + else: + from rixl._api import nixl_agent_config +except ImportError: + nixl_agent_config = None # type: ignore[assignment] + logger.warning_once("NIXL agent config is not available") + +try: + if current_platform.is_cuda(): + from nixl._bindings import nixlXferTelemetry + else: + from rixl._bindings import nixlXferTelemetry +except ImportError: + nixlXferTelemetry = None # type: ignore[assignment, misc] + +__all__ = ["NixlWrapper", "nixl_agent_config", "nixlXferTelemetry"] diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 1e05d68078e..443c0aae242 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4821,6 +4821,23 @@ class GPUModelRunner( ) self.model.set_aux_hidden_state_layers(aux_layers) + + if ( + is_mixture_of_experts(self.model) + and self.parallel_config.enable_eplb + and not load_dummy_weights + ): + logger.info_once( + "EPLB is enabled for model %s.", + self.model_config.model, + ) + assert self.eplb_state is not None + self.eplb_state.add_model( + self.model, + self.model_config, + ) + eplb_models += 1 + time_after_load = time.perf_counter() self.model_memory_usage = m.consumed_memory except torch.cuda.OutOfMemoryError as e: @@ -4860,15 +4877,10 @@ class GPUModelRunner( is_mixture_of_experts(self.model) and self.parallel_config.enable_eplb and not load_dummy_weights + and self.eplb_state is not None + and self.eplb_state.is_async ): - logger.info_once("EPLB is enabled for model %s.", self.model_config.model) - assert self.eplb_state is not None - self.eplb_state.add_model( - self.model, - self.model_config, - ) - if self.eplb_state.is_async: - self.eplb_state.start_async_loop() + self.eplb_state.start_async_loop() if ( self.vllm_config.compilation_config.mode From cc3993b05d008be370778466865cce1959ac85a3 Mon Sep 17 00:00:00 2001 From: zhanqiuhu <49648934+ZhanqiuHu@users.noreply.github.com> Date: Mon, 20 Apr 2026 06:39:08 -0400 Subject: [PATCH 139/150] nixl refactor [2/N]: unify TpKVTopology + HeteroTPTransferConfig into TransferTopology (#39529) Signed-off-by: Zhanqiu Hu --- .../unit/test_mooncake_connector.py | 8 +- .../kv_connector/unit/test_nixl_connector.py | 31 +- .../unit/test_nixl_connector_hma.py | 27 +- .../kv_transfer/kv_connector/utils.py | 1180 +++++++++-------- .../v1/mooncake/mooncake_connector.py | 22 +- .../kv_connector/v1/nixl/worker.py | 239 ++-- .../v1/ssm_conv_transfer_utils.py | 4 +- 7 files changed, 772 insertions(+), 739 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index 7b6fe3af0ce..83202e023c9 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -631,7 +631,7 @@ def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes(): mock_thread.return_value.is_alive.return_value = False worker.use_mla = True - worker.kv_topo.is_mla = True + worker.transfer_topo.is_mla = True # MLA cache tensor: shape[-2] is the block size. mla_cache = torch.zeros((2, 16, 96), dtype=torch.float16) @@ -692,9 +692,9 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size): # Override TP rank/size to simulate P TP=2 prefill_worker.tp_rank = P_TP_RANK prefill_worker.tp_size = P_TP_SIZE - # Update shared dict so kv_topo sees correct TP size prefill_worker._tp_size[prefill_worker.engine_id] = P_TP_SIZE - prefill_worker.kv_topo.tp_rank = P_TP_RANK + prefill_worker.transfer_topo.tp_rank = P_TP_RANK + prefill_worker.transfer_topo.tp_size = P_TP_SIZE prefill_worker.kv_caches_base_addr = [0x1000] prefill_worker.block_len_per_layer = [local_block_len] @@ -714,7 +714,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size): send_meta.ready.set() # Compute target D ranks using the production code path - target_d_ranks = prefill_worker.kv_topo.get_target_remote_ranks(d_tp_size) + target_d_ranks = prefill_worker.transfer_topo.handshake_target_ranks(d_tp_size) mock_socket = AsyncMock(spec=zmq.asyncio.Socket) mock_socket.send_multipart = AsyncMock() diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index d67b14e8dd4..50e83aa2ef2 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -21,7 +21,7 @@ from vllm import LLM from vllm.config import KVTransferConfig, set_current_vllm_config from vllm.distributed.kv_transfer.kv_connector.utils import ( KVOutputAggregator, - TpKVTopology, + TransferTopology, get_current_attn_backend, ) from vllm.distributed.kv_transfer.kv_connector.v1 import nixl @@ -463,19 +463,20 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): test_shape = self.attn_backends[0].get_kv_cache_shape( num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 ) - self.kv_topo = TpKVTopology( + self.transfer_topo = TransferTopology( tp_rank=self.tp_rank, + tp_size=self.world_size, + block_size=self.block_size, engine_id=self.engine_id, - remote_tp_size=self._tp_size, # shared state - remote_block_size=self._block_size, # shared state is_mla=self.use_mla, + is_mamba=False, total_num_kv_heads=self.model_config.get_total_num_kv_heads(), attn_backends=self.attn_backends, tensor_shape=test_shape, ) self.compat_hash = compute_nixl_compatibility_hash( - self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks + self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks ) def _nixl_handshake( @@ -496,7 +497,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): # Adjust remote block length metadata to satisfy heterogeneous TP # invariants enforced during handshake validation. remote_block_lens = list(self.block_len_per_layer) - tp_ratio = self.kv_topo.tp_ratio(remote_tp_size) + tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size) if remote_tp_size > self.world_size: # P TP > D TP case, block_len of remote is smaller remote_block_lens = [ @@ -731,8 +732,9 @@ class TestNixlHandshake: assert set(remote_agents.keys()) == set(range(tp_ratio)) remote_engine_id = worker.REMOTE_ENGINE_ID - assert worker._tp_size[remote_engine_id] == remote_tp_size - assert -tp_ratio == worker.kv_topo.tp_ratio_from_engine_id(remote_engine_id) + remote_info = worker.transfer_topo.get_engine_info(remote_engine_id) + assert remote_info.remote_tp_size == remote_tp_size + assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size) # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio assert len(worker.src_xfer_handles_by_tp_ratio[-tp_ratio]) == tp_ratio @@ -796,7 +798,7 @@ class TestNixlHandshake: (conn_p0.connector_worker, conn_p1.connector_worker) ): worker.world_size = p_tp_size - worker.kv_topo.remote_tp_size = {worker.engine_id: p_tp_size} + worker.transfer_topo.tp_size = p_tp_size worker.tp_rank = rank worker.use_mla = True @@ -2337,7 +2339,7 @@ def test_compatibility_hash_validation( remote_hash = compute_nixl_compatibility_hash( remote_vllm_config, decode_worker.backend_name, - decode_worker.kv_topo.cross_layers_blocks, + decode_worker.transfer_topo.cross_layers_blocks, ) prefill_block_size = config_overrides.get("block_size", 16) @@ -2424,12 +2426,13 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) test_shape = backend.get_kv_cache_shape( num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 ) - decode_worker.kv_topo = TpKVTopology( + decode_worker.transfer_topo = TransferTopology( tp_rank=decode_worker.tp_rank, + tp_size=decode_worker.world_size, + block_size=decode_worker.block_size, engine_id=decode_worker.engine_id, - remote_tp_size=decode_worker._tp_size, # shared state - remote_block_size=decode_worker._block_size, # shared state is_mla=decode_worker.use_mla, + is_mamba=False, total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(), attn_backends=[backend], tensor_shape=test_shape, @@ -2438,7 +2441,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) decode_worker.compat_hash = compute_nixl_compatibility_hash( decode_worker.vllm_config, decode_worker.backend_name, - decode_worker.kv_topo.cross_layers_blocks, + decode_worker.transfer_topo.cross_layers_blocks, ) if error_scenario == "handshake_decode_error": diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 30913ff98ee..5b609017359 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -152,13 +152,14 @@ def test_read_blocks_for_req_expands_remote_ids( remote_engine_id = "remote-engine" if has_mamba: - worker._mamba_phys_ratio = {remote_engine_id: remote_ratio} + worker._physical_blocks_per_logical = {remote_engine_id: remote_ratio} - # Mock kv_topo: empty remote ranks skips the transfer machinery entirely, - # isolating the block-ID expansion logic. - worker.kv_topo = MagicMock() - worker.kv_topo.get_target_remote_ranks_from_engine_id.return_value = [] - worker.kv_topo.tp_ratio_from_engine_id.return_value = 1 + # Mock transfer_topo: empty remote ranks skips the transfer machinery + # entirely, isolating the block-ID expansion logic. + worker.transfer_topo = MagicMock() + worker.transfer_topo.target_remote_ranks.return_value = [] + worker.transfer_topo.get_engine_info.return_value = MagicMock(remote_tp_size=1) + worker.transfer_topo.tp_ratio.return_value = 1 metadata = NixlConnectorMetadata() metadata.add_new_req_to_recv( @@ -317,7 +318,7 @@ def test_get_block_descs_ids_hybrid_ssm(): worker._has_mamba = True worker._is_mamba_group = [False, True] worker._physical_blocks_per_logical_kv_block = 1 - worker._mamba_phys_ratio = {engine_id: 1} + worker._physical_blocks_per_logical = {engine_id: 1} worker.block_len_per_layer = [100] # num_descs = num_regions * num_blocks (no blocks_first doubling) worker.num_descs = 2 * num_blocks @@ -355,7 +356,7 @@ def test_get_block_descs_ids_kernel_block_mismatch(): worker._has_mamba = True worker._is_mamba_group = [False, True] worker._physical_blocks_per_logical_kv_block = ratio - worker._mamba_phys_ratio = {engine_id: ratio} + worker._physical_blocks_per_logical = {engine_id: ratio} worker.block_len_per_layer = [100] worker.num_descs = 2 * num_blocks # 800 @@ -532,15 +533,15 @@ def test_has_mamba_init( ((9216, 524288), 4096, 131), ], ) -def test_compute_mamba_phys_ratio(ssm_sizes, block_len, expected_ratio): - """Verify that compute_mamba_phys_ratio is TP-dependent. +def test_compute_physical_blocks_per_logical(ssm_sizes, block_len, expected_ratio): + """Verify that compute_physical_blocks_per_logical is TP-dependent. With dimension-sharded Mamba state, the ratio differs across TP sizes (e.g. TP=1 → 261, TP=4 → 131 for Nemotron 30B). This is why - _mamba_phys_ratio must be stored per-engine. + _physical_blocks_per_logical must be stored per-engine. """ from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( - compute_mamba_phys_ratio, + compute_physical_blocks_per_logical, ) - assert compute_mamba_phys_ratio(ssm_sizes, block_len) == expected_ratio + assert compute_physical_blocks_per_logical(ssm_sizes, block_len) == expected_ratio diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index e75c1c0a3a4..63b56eddfae 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -5,7 +5,7 @@ KV cache helper for store. """ from collections.abc import Iterator -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, cast import torch @@ -319,535 +319,6 @@ def yield_req_data( ) -@dataclass -class TpKVTopology: - """ - Helper class for tensor parallel and KV topology information for - mapping between local and remote TP workers. - """ - - tp_rank: int - remote_tp_size: dict[EngineId, int] - is_mla: bool - total_num_kv_heads: int - attn_backends: list[type[AttentionBackend]] - engine_id: EngineId - remote_block_size: dict[EngineId, int] - tensor_shape: torch.Size | None = None - is_mamba: bool = False - - def __post_init__(self): - # Figure out whether the first dimension of the cache is K/V - # or num_blocks. This is used to register the memory regions correctly. - attn_backend = self.attn_backends[0] - if not self.is_mamba: - _MOCK_BLOCK_SIZE = 16 - kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape( - num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1 - ) - logger.debug("Test kv_cache_shape: %s", kv_cache_shape) - # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D], - # we just mock num_blocks to 1 for the dimension check below. - # Hybrid SSM models assume a single blocks_first layout - self._is_kv_layout_blocks_first = self.is_mamba or ( - len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1 - ) - - self._cross_layers_blocks = False - if self.tensor_shape is not None: - self._cross_layers_blocks = ( - len(self.tensor_shape) == len(kv_cache_shape) + 1 - ) - self.tensor_shape: torch.Size - - if self._cross_layers_blocks: - logger.debug("Using cross-layer KV cache") - # prepend layers dimension - _MOCK_NUM_LAYERS = 80 - kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape - try: - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order( - include_num_layers_dimension=self._cross_layers_blocks - ) - except (AttributeError, NotImplementedError): - assert self.tensor_shape is not None - kv_cache_stride_order = tuple(range(len(self.tensor_shape))) - - # In case of cross layers permute kv_cache_shape according to - # stride_order to retrieve physical position of block_size - kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) - - @property - def is_kv_layout_blocks_first(self) -> bool: - return self._is_kv_layout_blocks_first - - @property - def split_k_and_v(self) -> bool: - # Whether to register regions for K and V separately (when present). - return not ( - self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first - ) - - @property - def tp_size(self) -> int: - return self.remote_tp_size[self.engine_id] - - @property - def block_size(self) -> int: - return self.remote_block_size[self.engine_id] - - @property - def cross_layers_blocks(self) -> bool: - return self._cross_layers_blocks - - def tp_ratio( - self, - remote_tp_size: int, - ) -> int: - """ - Calculate the tensor parallel ratio between local and remote TP. - We can think of it as the number of local TP workers-per-remote TP - workers. Local workers will read from the same remote TP worker in - groups of size `tp_ratio`.If remote tp_size > local tp_size, the - ratio is flipped (remote_size/local_size) and the returned value is - negative. - """ - if self.tp_size >= remote_tp_size: - assert self.tp_size % remote_tp_size == 0, ( - f"Local tensor parallel size {self.tp_size} is not divisible " - f"by remote tensor parallel size {remote_tp_size}." - ) - return self.tp_size // remote_tp_size - - assert remote_tp_size % self.tp_size == 0, ( - f"Remote tensor parallel size {remote_tp_size} is not divisible " - f"by local tensor parallel size {self.tp_size}." - ) - # P TP > D TP case, return the ratio as negative - return -remote_tp_size // self.tp_size - - def block_size_ratio( - self, - remote_block_size: int, - ) -> int: - """ - Calculate the block size ratio between local and remote TP. - """ - assert self.block_size % remote_block_size == 0, ( - f"Local block size {self.block_size} is not divisible " - f"by remote block size {remote_block_size} or vice versa." - ) - return self.block_size // remote_block_size - - def tp_ratio_from_engine_id( - self, - remote_engine_id: EngineId, - ) -> int: - remote_tp_size = self.remote_tp_size[remote_engine_id] - return self.tp_ratio(remote_tp_size) - - def block_size_ratio_from_engine_id( - self, - remote_engine_id: EngineId, - ) -> int: - remote_block_size = self.remote_block_size[remote_engine_id] - return self.block_size_ratio(remote_block_size) - - def is_kv_replicated(self, engine_id: EngineId) -> bool: - """ - Whether the KV cache is replicated across TP workers due to the - number of TP workers being greater than the number of KV heads. - When they are equal, each TP rank still owns one distinct KV head, - so this is not considered replication. - """ - tp_size = self.remote_tp_size[engine_id] - return tp_size > self.total_num_kv_heads - - def replicates_kv_cache(self, remote_engine_id: EngineId) -> bool: - # MLA is always replicated as the hidden dim can't be split. - return self.is_mla or self.is_kv_replicated(remote_engine_id) - - def get_target_remote_ranks( - self, - remote_tp_size: int, - ) -> list[int]: - """ - Get the remote TP rank (on P) that the current local TP rank - (on D) will read from. When remote tp_size > local tp_size, we - read from multiple remote ranks. - """ - tp_ratio = self.tp_ratio(remote_tp_size) - if tp_ratio > 0: - return [self.tp_rank // tp_ratio] - - # P TP > D TP case, D reads from |tp_ratio| remote workers. - tp_ratio = -tp_ratio - return [self.tp_rank * tp_ratio + i for i in range(tp_ratio)] - - def get_target_remote_ranks_from_engine_id( - self, - remote_engine_id: EngineId, - ) -> list[int]: - remote_tp_size = self.remote_tp_size[remote_engine_id] - return self.get_target_remote_ranks(remote_tp_size) - - def get_transfer_cache_regions( - self, cache: torch.Tensor, layer_spec: "KVCacheSpec" - ) -> list[torch.Tensor] | torch.Tensor: - """Return the cache tensor(s) to register as NIXL memory regions, - also accounting for hybrid SSM models specificities. - """ - if isinstance(layer_spec, MambaSpec): - # Register the whole kv cache shared tensor, including SSM/Conv. This is - # similar to FI with the difference that SSM/Conv have different sizes - conv, ssm = cache - return [conv] - - # Check may be hacky but it's matching `_update_hybrid_attention_mamba_layout`. - if self.is_mamba and cache.shape[0] == 2: - # When MAMBA is present, all backends are blocks first, so that blocks - # can be shared between attention layers and mamba layers. Runner - # `_update_hybrid_attention_mamba_layout` already adjusted strides - # for FlashAttn-like backends so its num_blocks first. - # Swap [2<>num_blocks] dims to get required layout for hybrid SSM. - cache = cache.transpose(0, 1) - - # Regular case: backends like FA register K/V in separate regions - return cache if self.split_k_and_v else [cache] - - -# ---- Mamba-HMA hetero-TP transfer config ---- -# -# Key insight: with hetero-TP (P_TP > D_TP), FA KV cache may be -# replicated across P ranks (when P_TP > num_kv_heads), but Mamba -# conv/SSM state is almost always uniquely sharded per P rank. So the -# number of P ranks D must read from can differ between FA and Mamba, -# and they must be handled separately. - - -def _physical_head_range(tp_size: int, num_heads: int, rank: int) -> range: - """Physical KV head range stored in a rank's KV cache tensor. - - When ``tp_size <= num_heads``: sharded, K/TP contiguous heads per rank. - When ``tp_size > num_heads``: 1 physical head per rank. Heads are - distributed **contiguously** (matching vLLM's GQA weight partitioning): - consecutive ranks share a head before moving to the next one. - """ - if tp_size <= num_heads: - assert num_heads % tp_size == 0 - per_rank = num_heads // tp_size - return range(rank * per_rank, (rank + 1) * per_rank) - else: - h = rank * num_heads // tp_size - return range(h, h + 1) - - -def _range_overlap(a: range, b: range) -> range: - start = max(a.start, b.start) - stop = min(a.stop, b.stop) - return range(start, max(start, stop)) - - -@dataclass -class HeteroTPTransferConfig: - """Precomputed transfer plan for one (D rank, P engine) pair. - - Currently only instantiated for Mamba-HMA (hybrid SSM+Attention) models - where FA and mamba require different splitting factors. Could be extended - to other model types that need non-uniform hetero-TP transfer sizing. - - All descriptor sizes are computed here. The guarantee is: - local_entry_size == remote_entry_size (for NIXL) - - Attributes that start with ``fa_`` concern FlashAttention KV cache. - Attributes that start with ``mamba_`` concern Mamba conv/SSM state. - """ - - # ---- Input parameters (from handshake) ---- - tp_ratio: int - K: int # total_num_kv_heads (before TP sharding) - d_tp: int # D engine's tensor_parallel_size - p_tp: int # P engine's tensor_parallel_size - d_rank: int # this D worker's TP rank - use_mla: bool - - # Per-layer block lengths (bytes, K+V combined for blocks_first). - # Uniform across layers for current models. - d_block_len: int # D's block_len_per_layer (representative) - p_block_len: int # P's block_len_per_layer (from handshake) - is_blocks_first: bool # kv_topo.is_kv_layout_blocks_first - - # ---- Derived: computed in __post_init__ ---- - # - # Physical heads per rank (what the KV tensor actually stores) - d_physical_heads: int = field(init=False) - p_physical_heads: int = field(init=False) - - # How many distinct P ranks D needs for FA data - physical_fa_num_reads: int = field(init=False) - - # Which P ranks contribute unique FA heads (ordered by head index) - fa_read_targets: list[int] = field(init=False) - - # All P ranks needed for mamba (always abs_tp for tp_ratio < 0) - mamba_num_reads: int = field(init=False) - - # All P ranks this D rank communicates with (FA ∪ mamba) - transfer_targets: list[int] = field(init=False) - - # FA descriptor entry size (K or V side, for blocks_first layout) - # Guaranteed: fa_entry_size is the SAME for local handle AND remote desc. - fa_entry_size: int = field(init=False) - - # Replication flags - is_d_replicated: bool = field(init=False) - is_p_replicated: bool = field(init=False) - - # Pre-built set for fast lookup - _fa_target_set: frozenset[int] = field(init=False, repr=False) - # Map: P rank → index in fa_read_targets (for head slot offset) - _fa_target_index: dict[int, int] = field(init=False, repr=False) - - def __post_init__(self) -> None: - K = self.K - self.is_d_replicated = self.d_tp > K - self.is_p_replicated = self.p_tp > K - - self.d_physical_heads = max(1, K // self.d_tp) - self.p_physical_heads = max(1, K // self.p_tp) - - abs_tp = -self.tp_ratio if self.tp_ratio < 0 else 1 - - # ---- Mamba range (computed first so FA can prefer ranks in it) ---- - mamba_range: range | None = None - if self.tp_ratio < 0: - mamba_range = range(self.d_rank * abs_tp, (self.d_rank + 1) * abs_tp) - - # ---- FA read targets ---- - if self.use_mla or self.tp_ratio >= 0: - self.physical_fa_num_reads = 1 - self.fa_read_targets = ( - [0] - if self.use_mla - # Must match kv_topo.get_target_remote_ranks (d_rank // tp_ratio). - else [ - self.d_rank // self.tp_ratio if self.tp_ratio > 0 else self.d_rank - ] - ) - else: - d_needs = _physical_head_range(self.d_tp, K, self.d_rank) - # When mamba range exists, prefer P ranks within it so that - # FA targets are a subset of mamba transfer_targets (avoids - # orphaned FA targets outside the transfer loop). - search_range = mamba_range if mamba_range is not None else range(self.p_tp) - seen: set[tuple[int, int]] = set() - targets: list[int] = [] - for p in search_range: - p_has = _physical_head_range(self.p_tp, K, p) - ov = _range_overlap(d_needs, p_has) - if len(ov) > 0: - key = (ov.start, ov.stop) - if key not in seen: - seen.add(key) - targets.append(p) - if not targets: - # Fallback: search globally (should not happen in practice) - for p in range(self.p_tp): - p_has = _physical_head_range(self.p_tp, K, p) - ov = _range_overlap(d_needs, p_has) - if len(ov) > 0: - key = (ov.start, ov.stop) - if key not in seen: - seen.add(key) - targets.append(p) - self.fa_read_targets = targets - self.physical_fa_num_reads = len(targets) - - self._fa_target_set = frozenset(self.fa_read_targets) - self._fa_target_index = {r: i for i, r in enumerate(self.fa_read_targets)} - - # ---- Mamba targets ---- - if mamba_range is not None and abs_tp > self.physical_fa_num_reads: - self.mamba_num_reads = abs_tp - self.transfer_targets = list(mamba_range) - else: - self.mamba_num_reads = self.physical_fa_num_reads - self.transfer_targets = list(self.fa_read_targets) - - # ---- FA entry size ---- - # For blocks_first: block_len_per_layer includes K+V; // 2 gives K (or V). - # Use min(D, P) because D indexes into P when tp_ratio > 0, - # and P is the natural unit when tp_ratio < 0. - effective_block_len = min(self.d_block_len, self.p_block_len) - if self.is_blocks_first: - self.fa_entry_size = effective_block_len // 2 - else: - self.fa_entry_size = effective_block_len - - self._validate() - - def _validate(self) -> None: - """Cross-check internal consistency.""" - if self.is_d_replicated and self.is_p_replicated and self.tp_ratio > 0: - logger.info( - "Both-replicated hetero-TP: D_TP=%d > P_TP=%d > K=%d. " - "Using d_rank // tp_ratio routing with relative head offset.", - self.d_tp, - self.p_tp, - self.K, - ) - - # FA targets must be a subset of transfer_targets - tt_set = set(self.transfer_targets) - for t in self.fa_read_targets: - if t not in tt_set: - logger.error( - "FA target P rank %d is NOT in transfer_targets %s. " - "This will cause missed FA reads!", - t, - self.transfer_targets, - ) - - # For tp_ratio < 0 with blocks_first: D_K_half / reads should == P_K_half - if ( - self.is_blocks_first - and self.tp_ratio < 0 - and self.physical_fa_num_reads > 0 - ): - d_k_half = self.d_block_len // 2 - p_k_half = self.p_block_len // 2 - expected_local = d_k_half // self.physical_fa_num_reads - if expected_local != p_k_half: - logger.warning( - "FA size mismatch: D_K_half=%d / reads=%d = %d, " - "but P_K_half=%d. This may indicate a head count or " - "Mamba-HMA inflation inconsistency.", - d_k_half, - self.physical_fa_num_reads, - expected_local, - p_k_half, - ) - - # ---- Query methods ---- - - def should_skip_fa(self, p_rank: int) -> bool: - """Whether to skip FA groups for this P rank (mamba-only transfer).""" - return p_rank not in self._fa_target_set - - def fa_head_slot(self, p_rank: int) -> int: - """Index into D's FA block for this P rank's head data. - - For P ranks in fa_read_targets, returns 0, 1, ..., reads-1. - For P ranks NOT in fa_read_targets (replicated duplicates), - returns the slot of the matching FA target with the same head. - """ - if p_rank in self._fa_target_index: - return self._fa_target_index[p_rank] - # Duplicate head: find which fa_target has the same physical head - p_head = _physical_head_range(self.p_tp, self.K, p_rank) - for target in self.fa_read_targets: - t_head = _physical_head_range(self.p_tp, self.K, target) - if _range_overlap(p_head, t_head): - return self._fa_target_index[target] - return 0 # fallback - - def fa_rank_offset(self, remote_kv_block_len: int) -> int: - """Byte offset into P's FA block for this D rank. - - When D is replicated (D_TP > K), multiple D ranks share a head. - Computes offset *relative to the target P rank's first head* - so it works regardless of how many heads P has. - When neither side replicates, falls back to tp_rank % tp_ratio. - Returns 0 when D does not index into P's block. - """ - if self.use_mla or self.tp_ratio <= 0: - return 0 - if self.is_d_replicated: - d_head = self.d_rank * self.K // self.d_tp - p_rank = self.fa_read_targets[0] - p_start = p_rank * self.K // self.p_tp - return (d_head - p_start) * remote_kv_block_len - return self.d_rank % self.tp_ratio * remote_kv_block_len - - @property - def needs_split_handles(self) -> bool: - """Whether per-P-rank split handles are needed. - - True when FA and mamba have different read counts, requiring - different splitting factors in the local handle. - """ - return self.tp_ratio < 0 and not self.use_mla and len(self.transfer_targets) > 1 - - def compute_split_handle_data( - self, - src_blocks_data: list[tuple[int, int, int]], - num_fa_descs: int, - abs_tp: int, - ) -> list[list[tuple[int, int, int]]]: - """Compute per-P-rank (addr, len, tp) triples for Mamba-HMA split handles. - - FA descriptors (indices < num_fa_descs) are sliced by - ``physical_fa_num_reads``; mamba descriptors are sliced uniformly - by ``abs_tp``. - - Returns one list of triples per transfer target. - """ - all_handle_data: list[list[tuple[int, int, int]]] = [] - for p_idx, p_rank in enumerate(self.transfer_targets): - handle_data: list[tuple[int, int, int]] = [] - skip_fa = self.should_skip_fa(p_rank) - fa_slot = self.fa_head_slot(p_rank) if not skip_fa else 0 - - for j, (addr, local_len, tp) in enumerate(src_blocks_data): - if j < num_fa_descs: - assert self.physical_fa_num_reads >= 1 - fa_chunk = local_len // self.physical_fa_num_reads - handle_data.append((addr + fa_slot * fa_chunk, fa_chunk, tp)) - else: - mamba_chunk = local_len // abs_tp - handle_data.append((addr + p_idx * mamba_chunk, mamba_chunk, tp)) - all_handle_data.append(handle_data) - return all_handle_data - - def filter_block_ids_for_rank( - self, - remote_rank: int, - local_ids: BlockIds, - remote_ids: BlockIds, - is_mamba_group: list[bool], - ) -> tuple[BlockIds, BlockIds]: - """Zero out FA groups for P ranks outside fa_read_targets. - - Returns (filtered_local_ids, filtered_remote_ids). When the - remote rank carries FA data for this D rank, returns the inputs - unchanged. - """ - if not self.should_skip_fa(remote_rank): - return local_ids, remote_ids - num_groups = len(local_ids) - filtered_local: list[list[int]] = [ - [] if not is_mamba_group[g] else local_ids[g] for g in range(num_groups) - ] - filtered_remote: list[list[int]] = [ - [] if not is_mamba_group[g] else remote_ids[g] for g in range(num_groups) - ] - return filtered_local, filtered_remote - - def describe(self) -> str: - """One-line summary for logging.""" - return ( - f"HeteroTPTransferConfig(" - f"tp_ratio={self.tp_ratio}, K={self.K}, " - f"d_tp={self.d_tp}, p_tp={self.p_tp}, d_rank={self.d_rank}, " - f"physical_fa_reads={self.physical_fa_num_reads}, " - f"mamba_reads={self.mamba_num_reads}, " - f"fa_targets={self.fa_read_targets}, " - f"transfer_targets={self.transfer_targets}, " - f"fa_entry_size={self.fa_entry_size}, " - f"d_block_len={self.d_block_len}, p_block_len={self.p_block_len})" - ) - - def get_current_attn_backends( vllm_config: VllmConfig, layer_names: list[str] | None = None ) -> list[type[AttentionBackend]]: @@ -893,48 +364,607 @@ def get_current_attn_backend( return get_current_attn_backends(vllm_config, layer_names)[0] -# TODO (ZhanqiuHu): Consolidate TpKVTopology and HeteroTPTransferConfig -# into a single engine-agnostic TransferTopology class. -# 6 of 9 HeteroTPTransferConfig init fields duplicate TpKVTopology data. -# -# @dataclass -# class EngineTransferInfo: -# """Per-remote-engine transfer state, computed at handshake.""" -# p_tp: int -# tp_ratio: int -# p_block_len: int -# block_size: int -# # Mamba-specific (None for non-mamba models) -# fa_read_targets: list[int] | None = None -# transfer_targets: list[int] | None = None -# physical_fa_num_reads: int | None = None -# mamba_num_reads: int | None = None -# fa_entry_size: int | None = None -# -# class TransferTopology: -# """Single source of truth for TP topology + transfer sizing.""" -# # Shared (set once at init, replaces duplicate fields) -# tp_rank: int # == TpKVTopology.tp_rank == HeteroTP.d_rank -# tp_size: int # == TpKVTopology.tp_size == HeteroTP.d_tp -# total_num_kv_heads: int # == HeteroTP.K -# is_mla: bool # == HeteroTP.use_mla -# is_mamba: bool -# is_blocks_first: bool # == HeteroTP.is_blocks_first -# d_block_len: int -# -# # Per-engine (populated via register_engine() at handshake) -# _engines: dict[EngineId, EngineTransferInfo] -# -# def register_engine(self, engine_id, p_tp, p_block_len, ...): ... -# -# # General (from TpKVTopology) -# def tp_ratio(self, engine_id) -> int: ... -# def target_remote_ranks(self, engine_id) -> list[int]: ... -# def is_kv_replicated(self, engine_id) -> bool: ... -# -# # Mamba-specific (from HeteroTPTransferConfig, gated by is_mamba) -# def fa_rank_offset(self, engine_id, block_len) -> int: ... -# def physical_fa_num_reads(self, engine_id) -> int: ... -# def transfer_targets(self, engine_id) -> list[int]: ... -# def should_skip_fa(self, engine_id, p_rank) -> bool: ... -# def filter_block_ids_for_rank(self, engine_id, ...) -> ...: ... +# ---- Per-engine transfer info ---- + + +@dataclass(frozen=True) +class EngineTransferInfo: + """Common per-remote-engine transfer state, computed at handshake. + + Stored per ``engine_id`` inside ``TransferTopology._engines``. + """ + + remote_tp_size: int + + remote_block_len: int + """Block length (bytes)""" + + remote_block_size: int + """Tokens per block.""" + + remote_physical_blocks_per_logical: int + """Physical blocks per logical block.""" + + +@dataclass(frozen=True) +class MambaEngineTransferInfo(EngineTransferInfo): + """Extends ``EngineTransferInfo`` with Mamba-hybrid transfer geometry. + + For hybrid SSM+Attention models, FA and Mamba layers may require + different numbers of reads from different remote ranks. This + dataclass captures that per-engine transfer plan. + """ + + remote_fa_source_ranks: tuple[int, ...] + """Remote ranks carrying unique FA heads for this local rank.""" + + remote_all_source_ranks: tuple[int, ...] + """All remote ranks this local rank reads from (FA + Mamba).""" + + remote_num_fa_reads: int + """Number of distinct remote ranks needed for FA data.""" + + remote_num_mamba_reads: int + """Number of distinct remote ranks needed for Mamba data.""" + + remote_fa_descriptor_bytes: int + """Byte size of one FA K (or V) descriptor entry.""" + + is_remote_replicated: bool + """Whether the remote engine has replicated KV heads + (remote_tp_size > total_num_kv_heads).""" + + remote_physical_heads: int + """Physical KV heads stored per remote rank.""" + + +# ---- Transfer topology ---- + + +@dataclass +class TransferTopology: + """Single source of truth for local TP identity and per-engine remote info.""" + + tp_rank: int + tp_size: int + block_size: int + engine_id: EngineId + is_mla: bool + is_mamba: bool + total_num_kv_heads: int + attn_backends: list[type[AttentionBackend]] + tensor_shape: torch.Size | None = None + + def __post_init__(self): + self.local_physical_heads = max(1, self.total_num_kv_heads // self.tp_size) + + self._engines: dict[EngineId, EngineTransferInfo] = {} + self._fa_source_sets: dict[EngineId, frozenset[int]] = {} + self._fa_source_indices: dict[EngineId, dict[int, int]] = {} + + # Figure out whether the first dimension of the cache is K/V + # or num_blocks. + attn_backend = self.attn_backends[0] + if not self.is_mamba: + _MOCK_BLOCK_SIZE = 16 + kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape( + num_blocks=1, + block_size=_MOCK_BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + ) + logger.debug("Test kv_cache_shape: %s", kv_cache_shape) + # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D], + # we just mock num_blocks to 1 for the dimension check below. + # Hybrid SSM models assume a single blocks_first layout + self._is_kv_layout_blocks_first = self.is_mamba or ( + len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1 + ) + + self._cross_layers_blocks = False + if self.tensor_shape is not None: + self._cross_layers_blocks = ( + len(self.tensor_shape) == len(kv_cache_shape) + 1 + ) + + if self._cross_layers_blocks: + logger.debug("Using cross-layer KV cache") + _MOCK_NUM_LAYERS = 80 + kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape + try: + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order( + include_num_layers_dimension=self._cross_layers_blocks + ) + except (AttributeError, NotImplementedError): + assert self.tensor_shape is not None + kv_cache_stride_order = tuple(range(len(self.tensor_shape))) + kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) + + # ============================================================ + # Engine registration + # ============================================================ + + def register_remote_engine( + self, + remote_engine_id: EngineId, + remote_tp_size: int, + remote_block_size: int, + remote_block_len: int, + remote_physical_blocks_per_logical: int, + *, + local_block_len: int = 0, + ) -> EngineTransferInfo: + """Register a remote engine, unifying worker dicts state. + + Only remote engines should be registered here — the local engine's + identity (tp_size, block_size, etc.) is set via ``__init__`` params. + + For Mamba models, also computes the Mamba transfer plan and + builds the FA source lookup caches. + + Args: + local_block_len: Local representative block_len (bytes). + Required for Mamba models to compute ``fa_descriptor_bytes``. + """ + assert remote_engine_id != self.engine_id, ( + f"Cannot register local engine {self.engine_id} as remote. " + f"Local identity is set via __init__ params." + ) + if remote_engine_id in self._engines: + return self._engines[remote_engine_id] + info: EngineTransferInfo + if self.is_mamba: + info = self._build_mamba_info( + remote_tp_size=remote_tp_size, + remote_block_size=remote_block_size, + remote_block_len=remote_block_len, + remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical), + local_block_len=local_block_len, + ) + assert isinstance(info, MambaEngineTransferInfo) + self._fa_source_sets[remote_engine_id] = frozenset( + info.remote_fa_source_ranks + ) + self._fa_source_indices[remote_engine_id] = { + r: i for i, r in enumerate(info.remote_fa_source_ranks) + } + else: + info = EngineTransferInfo( + remote_tp_size=remote_tp_size, + remote_block_len=remote_block_len, + remote_block_size=remote_block_size, + remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical), + ) + self._engines[remote_engine_id] = info + return info + + def get_engine_info(self, remote_engine_id: EngineId) -> EngineTransferInfo: + return self._engines[remote_engine_id] + + # ============================================================ + # Layout properties + # ============================================================ + + @property + def is_kv_layout_blocks_first(self) -> bool: + return self._is_kv_layout_blocks_first + + @property + def cross_layers_blocks(self) -> bool: + return self._cross_layers_blocks + + @property + def split_k_and_v(self) -> bool: + # Whether to register regions for K and V separately (when present). + return not ( + self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first + ) + + # ============================================================ + # Common methods + # ============================================================ + + def tp_ratio(self, remote_tp_size: int) -> int: + """Calculate the tensor parallel ratio between local and remote TP. + + Positive when local_tp >= remote_tp (local workers read from the + same remote worker in groups of size ``tp_ratio``). Negative when + remote_tp > local_tp (ratio is flipped). + """ + if self.tp_size >= remote_tp_size: + assert self.tp_size % remote_tp_size == 0, ( + f"Local tensor parallel size {self.tp_size} is not divisible " + f"by remote tensor parallel size {remote_tp_size}." + ) + return self.tp_size // remote_tp_size + assert remote_tp_size % self.tp_size == 0, ( + f"Remote tensor parallel size {remote_tp_size} is not divisible " + f"by local tensor parallel size {self.tp_size}." + ) + return -(remote_tp_size // self.tp_size) + + def block_size_ratio(self, remote_block_size: int) -> int: + """Calculate the block size ratio between local and remote.""" + assert self.block_size % remote_block_size == 0, ( + f"Local block size {self.block_size} is not divisible " + f"by remote block size {remote_block_size} or vice versa." + ) + return self.block_size // remote_block_size + + def is_kv_replicated(self, remote_engine_id: EngineId) -> bool: + """Whether the KV cache is replicated across TP workers due to the + number of TP workers being greater than the number of KV heads. + """ + return self._engines[remote_engine_id].remote_tp_size > self.total_num_kv_heads + + def replicates_kv_cache(self, remote_engine_id: EngineId) -> bool: + # MLA is always replicated as the hidden dim can't be split. + return self.is_mla or self.is_kv_replicated(remote_engine_id) + + @property + def local_replicates_kv_cache(self) -> bool: + """Whether the local engine's KV cache is replicated.""" + return self.is_mla or self.tp_size > self.total_num_kv_heads + + def handshake_target_ranks(self, remote_tp_size: int) -> list[int]: + """Pre-registration: compute which remote TP ranks to handshake with. + + Pure math based on local/remote TP sizes — does not require + the remote engine to be registered yet. + """ + tp_ratio = self.tp_ratio(remote_tp_size) + if tp_ratio > 0: + return [self.tp_rank // tp_ratio] + abs_ratio = -tp_ratio + return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)] + + def target_remote_ranks(self, remote_engine_id: EngineId) -> list[int]: + """Get the remote TP rank(s) that the current local TP rank will + read from. When remote tp_size > local tp_size, reads from + multiple remote ranks. + + For Mamba models, returns the precomputed ``all_source_ranks`` + (FA + Mamba union). + """ + info = self._engines[remote_engine_id] + if isinstance(info, MambaEngineTransferInfo): + return list(info.remote_all_source_ranks) + + tp_ratio = self.tp_ratio(info.remote_tp_size) + if tp_ratio > 0: + return [self.tp_rank // tp_ratio] + # remote TP > local TP: read from |tp_ratio| remote workers + abs_ratio = -tp_ratio + return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)] + + def get_transfer_cache_regions( + self, cache: torch.Tensor, layer_spec: "KVCacheSpec" + ) -> list[torch.Tensor] | torch.Tensor: + """Return the cache tensor(s) to register as NIXL memory regions, + also accounting for hybrid SSM models specificities. + """ + if isinstance(layer_spec, MambaSpec): + # Register the whole kv cache shared tensor, including + # SSM/Conv. + conv, ssm = cache + return [conv] + + # Check may be hacky but it's matching + # `_update_hybrid_attention_mamba_layout`. + if self.is_mamba and cache.shape[0] == 2: + # When MAMBA is present, all backends are blocks first, so + # that blocks can be shared between attention layers and mamba + # layers. Runner already adjusted strides for FlashAttn-like + # backends so its num_blocks first. + # Swap [2<>num_blocks] dims for hybrid SSM layout. + cache = cache.transpose(0, 1) + + # Regular case: backends like FA register K/V in separate regions + return cache if self.split_k_and_v else [cache] + + # ============================================================ + # Mamba-specific methods + # ============================================================ + + def should_skip_fa(self, remote_engine_id: EngineId, remote_rank: int) -> bool: + """Whether to skip FA groups for this remote rank (mamba-only).""" + return remote_rank not in self._fa_source_sets[remote_engine_id] + + def fa_head_slot(self, remote_engine_id: EngineId, remote_rank: int) -> int: + """Index into local FA block for this remote rank's head data. + + For remote ranks in ``fa_source_ranks``, returns 0, 1, …, reads-1. + For ranks NOT in ``fa_source_ranks`` (replicated duplicates), + returns the slot of the matching source rank with the same head. + """ + fa_index = self._fa_source_indices[remote_engine_id] + if remote_rank in fa_index: + return fa_index[remote_rank] + mamba_info = self._engines[remote_engine_id] + assert isinstance(mamba_info, MambaEngineTransferInfo) + K = self.total_num_kv_heads + remote_tp = mamba_info.remote_tp_size + r_head = self._physical_head_range(remote_tp, K, remote_rank) + for target in mamba_info.remote_fa_source_ranks: + t_head = self._physical_head_range(remote_tp, K, target) + if self._range_overlap(r_head, t_head): + return fa_index[target] + return 0 + + def fa_rank_offset( + self, remote_engine_id: EngineId, remote_kv_block_len: int + ) -> int: + """Byte offset into remote FA block for this local rank. + + When local TP is replicated (local_tp > K), multiple local ranks + share a head. Computes offset *relative to the target remote + rank's first head* so it works regardless of how many heads the + remote has. Returns 0 when local does not index into remote. + """ + mamba_info = self._engines[remote_engine_id] + assert isinstance(mamba_info, MambaEngineTransferInfo) + tp_ratio = self.tp_ratio(mamba_info.remote_tp_size) + if self.is_mla or tp_ratio <= 0: + return 0 + K = self.total_num_kv_heads + is_local_replicated = self.tp_size > K + if is_local_replicated: + local_head = self.tp_rank * K // self.tp_size + p_rank = mamba_info.remote_fa_source_ranks[0] + p_start = p_rank * K // mamba_info.remote_tp_size + return (local_head - p_start) * remote_kv_block_len + return self.tp_rank % tp_ratio * remote_kv_block_len + + def needs_split_handles(self, remote_engine_id: EngineId) -> bool: + """Whether per-remote-rank split handles are needed. + + True when FA and mamba have different read counts, requiring + different splitting factors in the local handle. + """ + mamba_info = self._engines[remote_engine_id] + assert isinstance(mamba_info, MambaEngineTransferInfo) + tp_ratio = self.tp_ratio(mamba_info.remote_tp_size) + return ( + tp_ratio < 0 + and not self.is_mla + and len(mamba_info.remote_all_source_ranks) > 1 + ) + + def compute_split_handle_data( + self, + remote_engine_id: EngineId, + src_blocks_data: list[tuple[int, int, int]], + num_fa_descs: int, + abs_tp: int, + ) -> list[list[tuple[int, int, int]]]: + """Per-remote-rank (addr, len, dev) triples for Mamba-HMA split + handles. + + FA descriptors (indices < num_fa_descs) are sliced by + ``remote_num_fa_reads``; mamba descriptors are sliced uniformly + by ``abs_tp``. + """ + mamba_info = self._engines[remote_engine_id] + assert isinstance(mamba_info, MambaEngineTransferInfo) + all_handle_data: list[list[tuple[int, int, int]]] = [] + for p_idx, p_rank in enumerate(mamba_info.remote_all_source_ranks): + handle_data: list[tuple[int, int, int]] = [] + skip_fa = self.should_skip_fa(remote_engine_id, p_rank) + fa_slot = self.fa_head_slot(remote_engine_id, p_rank) if not skip_fa else 0 + for j, (addr, local_len, dev) in enumerate(src_blocks_data): + if j < num_fa_descs: + assert mamba_info.remote_num_fa_reads >= 1 + fa_chunk = local_len // mamba_info.remote_num_fa_reads + handle_data.append((addr + fa_slot * fa_chunk, fa_chunk, dev)) + else: + mamba_chunk = local_len // abs_tp + handle_data.append((addr + p_idx * mamba_chunk, mamba_chunk, dev)) + all_handle_data.append(handle_data) + return all_handle_data + + def filter_block_ids_for_rank( + self, + remote_engine_id: EngineId, + remote_rank: int, + local_ids: BlockIds, + remote_ids: BlockIds, + is_mamba_group: list[bool], + ) -> tuple[BlockIds, BlockIds]: + """Zero out FA groups for remote ranks outside ``fa_source_ranks``. + + Returns (filtered_local_ids, filtered_remote_ids). When the + remote rank carries FA data for this local rank, returns the + inputs unchanged. + """ + if not self.should_skip_fa(remote_engine_id, remote_rank): + return local_ids, remote_ids + num_groups = len(local_ids) + filtered_local: list[list[int]] = [ + [] if not is_mamba_group[g] else local_ids[g] for g in range(num_groups) + ] + filtered_remote: list[list[int]] = [ + [] if not is_mamba_group[g] else remote_ids[g] for g in range(num_groups) + ] + return filtered_local, filtered_remote + + def describe(self, remote_engine_id: EngineId) -> str: + """One-line summary of transfer config for logging.""" + info = self._engines[remote_engine_id] + base = ( + f"tp_ratio={self.tp_ratio(info.remote_tp_size)}, " + f"K={self.total_num_kv_heads}, " + f"local_tp={self.tp_size}, " + f"remote_tp={info.remote_tp_size}, " + f"local_rank={self.tp_rank}, " + f"remote_block_len={info.remote_block_len}" + ) + if isinstance(info, MambaEngineTransferInfo): + return ( + f"TransferTopology.mamba({base}, " + f"fa_reads={info.remote_num_fa_reads}, " + f"mamba_reads={info.remote_num_mamba_reads}, " + f"fa_sources={list(info.remote_fa_source_ranks)}, " + f"all_sources={list(info.remote_all_source_ranks)}, " + f"fa_desc_bytes={info.remote_fa_descriptor_bytes})" + ) + return f"TransferTopology({base})" + + # ============================================================ + # Private helpers + # ============================================================ + # Mamba-HMA hetero-TP transfer config: + # With hetero-TP (P_TP > D_TP), FA KV cache may be replicated across + # P ranks (when P_TP > num_kv_heads), but Mamba conv/SSM state is + # almost always uniquely sharded per P rank. So the number of P + # ranks D must read from can differ between FA and Mamba, and they + # must be handled separately. + + @staticmethod + def _physical_head_range(tp_size: int, num_heads: int, rank: int) -> range: + """Physical KV head range stored in a rank's KV cache tensor. + + When ``tp_size <= num_heads``: sharded, K/TP contiguous heads per rank. + When ``tp_size > num_heads``: 1 physical head per rank. Heads are + distributed **contiguously** (matching vLLM's GQA weight partitioning): + consecutive ranks share a head before moving to the next one. + """ + if tp_size <= num_heads: + assert num_heads % tp_size == 0 + per_rank = num_heads // tp_size + return range(rank * per_rank, (rank + 1) * per_rank) + else: + h = rank * num_heads // tp_size + return range(h, h + 1) + + @staticmethod + def _range_overlap(a: range, b: range) -> range: + start = max(a.start, b.start) + stop = min(a.stop, b.stop) + return range(start, max(start, stop)) + + # ============================================================ + # Private: build Mamba transfer info + # ============================================================ + + def _build_mamba_info( + self, + remote_tp_size: int, + remote_block_size: int, + remote_block_len: int, + remote_physical_blocks_per_logical: int, + local_block_len: int, + ) -> MambaEngineTransferInfo: + """Compute Mamba transfer plan.""" + K = self.total_num_kv_heads + local_tp = self.tp_size + local_rank = self.tp_rank + + is_remote_replicated = remote_tp_size > K + remote_physical_heads = max(1, K // remote_tp_size) + + if local_tp >= remote_tp_size: + assert local_tp % remote_tp_size == 0 + tp_ratio = local_tp // remote_tp_size + else: + assert remote_tp_size % local_tp == 0 + tp_ratio = -(remote_tp_size // local_tp) + + abs_tp = -tp_ratio if tp_ratio < 0 else 1 + + mamba_range: range | None = None + if tp_ratio < 0: + mamba_range = range(local_rank * abs_tp, (local_rank + 1) * abs_tp) + + # ---- FA read targets ---- + if self.is_mla or tp_ratio >= 0: + num_fa_reads = 1 + fa_source_ranks: list[int] = ( + [0] + if self.is_mla + else [local_rank // tp_ratio if tp_ratio > 0 else local_rank] + ) + else: + local_needs = self._physical_head_range(local_tp, K, local_rank) + search_range = ( + mamba_range if mamba_range is not None else range(remote_tp_size) + ) + seen: set[tuple[int, int]] = set() + fa_source_ranks = [] + for p in search_range: + p_has = self._physical_head_range(remote_tp_size, K, p) + ov = self._range_overlap(local_needs, p_has) + if len(ov) > 0: + key = (ov.start, ov.stop) + if key not in seen: + seen.add(key) + fa_source_ranks.append(p) + if not fa_source_ranks: + for p in range(remote_tp_size): + p_has = self._physical_head_range(remote_tp_size, K, p) + ov = self._range_overlap(local_needs, p_has) + if len(ov) > 0: + key = (ov.start, ov.stop) + if key not in seen: + seen.add(key) + fa_source_ranks.append(p) + num_fa_reads = len(fa_source_ranks) + + # ---- All source ranks (mamba + FA) ---- + if mamba_range is not None and abs_tp > num_fa_reads: + num_mamba_reads = abs_tp + all_source_ranks = list(mamba_range) + else: + num_mamba_reads = num_fa_reads + all_source_ranks = list(fa_source_ranks) + + # ---- FA descriptor bytes ---- + effective_block_len = min(local_block_len, remote_block_len) + if self.is_kv_layout_blocks_first: + fa_descriptor_bytes = effective_block_len // 2 + else: + fa_descriptor_bytes = effective_block_len + + # ---- Validation ---- + is_local_replicated = local_tp > K + if is_local_replicated and is_remote_replicated and tp_ratio > 0: + logger.info( + "Both-replicated hetero-TP: local_tp=%d > remote_tp=%d > K=%d.", + local_tp, + remote_tp_size, + K, + ) + tt_set = set(all_source_ranks) + for t in fa_source_ranks: + if t not in tt_set: + logger.error( + "FA source rank %d NOT in all_source_ranks %s.", + t, + all_source_ranks, + ) + if self.is_kv_layout_blocks_first and tp_ratio < 0 and num_fa_reads > 0: + local_k_half = local_block_len // 2 + remote_k_half = remote_block_len // 2 + expected = local_k_half // num_fa_reads + if expected != remote_k_half: + logger.warning( + "FA size mismatch: local_k_half=%d / reads=%d = %d, " + "but remote_k_half=%d.", + local_k_half, + num_fa_reads, + expected, + remote_k_half, + ) + + return MambaEngineTransferInfo( + remote_tp_size=remote_tp_size, + remote_block_len=remote_block_len, + remote_block_size=remote_block_size, + remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical), + remote_fa_source_ranks=tuple(fa_source_ranks), + remote_all_source_ranks=tuple(all_source_ranks), + remote_num_fa_reads=num_fa_reads, + remote_num_mamba_reads=num_mamba_reads, + remote_fa_descriptor_bytes=fa_descriptor_bytes, + is_remote_replicated=is_remote_replicated, + remote_physical_heads=remote_physical_heads, + ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 67603e10ff6..2057c79fa58 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -21,7 +21,7 @@ from vllm import envs from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.utils import ( EngineId, - TpKVTopology, + TransferTopology, get_current_attn_backend, get_current_attn_backends, ) @@ -764,13 +764,13 @@ class MooncakeConnectorWorker: logger.debug("Detected kv cache layout %s", self.kv_cache_layout) self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size} - self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size} - self.kv_topo = TpKVTopology( + self.transfer_topo = TransferTopology( tp_rank=self.tp_rank, + tp_size=self.tp_size, + block_size=self.block_size, engine_id=self.engine_id, - remote_tp_size=self._tp_size, # shared state - remote_block_size=self._block_size, # shared state is_mla=self.use_mla, + is_mamba=False, total_num_kv_heads=self.model_config.get_total_num_kv_heads(), attn_backends=[backend], ) @@ -911,7 +911,7 @@ class MooncakeConnectorWorker: self, identity: bytes, sock: zmq.asyncio.Socket, meta: MooncakeXferMetadata ): pending_reqs: dict[ReqId, SendBlockMeta] = {} - remote_tp_ranks = self.kv_topo.get_target_remote_ranks(meta.remote_tp_size) + remote_tp_ranks = self.transfer_topo.handshake_target_ranks(meta.remote_tp_size) if meta.remote_tp_rank not in remote_tp_ranks: # This D worker does not pair with the P worker. msg = ( @@ -1256,7 +1256,7 @@ class MooncakeConnectorWorker: seen_base_addresses = [] self.block_len_per_layer = [] - split_k_and_v = self.kv_topo.split_k_and_v + split_k_and_v = self.transfer_topo.split_k_and_v tensor_size_bytes = None for layer_name, cache_or_caches in kv_caches.items(): cache_list = cache_or_caches if split_k_and_v else [cache_or_caches] @@ -1495,8 +1495,8 @@ class MooncakeConnectorWorker: remote_engine_id: EngineId, pull_metas: dict[ReqId, PullReqMeta], ): - remote_tp_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id( - remote_engine_id + remote_tp_ranks = self.transfer_topo.handshake_target_ranks( + self._tp_size[remote_engine_id] ) count = len(remote_tp_ranks) logger.debug( @@ -1587,7 +1587,7 @@ class MooncakeConnectorWorker: ) def _producer_cache_is_replicated(self) -> bool: - return self.kv_topo.replicates_kv_cache(self.engine_id) + return self.transfer_topo.local_replicates_kv_cache def _get_transfer_regions( self, base_addrs: list[int], block_lens: list[int] @@ -1595,7 +1595,7 @@ class MooncakeConnectorWorker: return _expand_transfer_regions( base_addrs=base_addrs, block_lens=block_lens, - is_kv_layout_blocks_first=self.kv_topo.is_kv_layout_blocks_first, + is_kv_layout_blocks_first=self.transfer_topo.is_kv_layout_blocks_first, ) def _get_sender_transfer_plan( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py index 8157df63d1b..bd7ef5973f6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py @@ -21,8 +21,8 @@ from vllm import envs from vllm.distributed.kv_transfer.kv_connector.utils import ( BlockIds, EngineId, - HeteroTPTransferConfig, - TpKVTopology, + MambaEngineTransferInfo, + TransferTopology, get_current_attn_backends, kv_postprocess_blksize_and_layout_on_receive, kv_postprocess_blksize_on_receive, @@ -49,7 +49,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( ) from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( MambaConvSplitInfo, - compute_mamba_phys_ratio, + compute_physical_blocks_per_logical, derive_mamba_conv_split, ) from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config @@ -269,14 +269,12 @@ class NixlConnectorWorker: self._registered_descs: list[Any] = [] # ---- Mamba-HMA per-engine state (only used when self._has_mamba) ---- - # Per-engine transfer config (source of truth for FA/mamba sizing). - self._transfer_configs: dict[str, HeteroTPTransferConfig] = {} - # NOTE (ZhanqiuHu): _mamba_phys_ratio MUST be per-engine. - # compute_mamba_phys_ratio = ceil((conv_bytes + ssm_bytes) / block_len) + # NOTE (ZhanqiuHu): _physical_blocks_per_logical MUST be per-engine. + # physical_blocks_per_logical = ceil((conv_bytes + ssm_bytes) / block_len) # where conv/ssm bytes are per-TP-rank (dimension-sharded). With # heterogeneous TP the per-rank sizes differ, so the ratio differs: # e.g. Nemotron 30B: P(TP=4) → 131, D(TP=1) → 261. - self._mamba_phys_ratio: dict[EngineId, int] = {} + self._physical_blocks_per_logical: dict[EngineId, int] = {} # In progress transfers. # [req_id -> list[handle]] @@ -322,10 +320,8 @@ class NixlConnectorWorker: # lazy initialized in register_kv_caches self.compat_hash: str | None = None - self.kv_topo: TpKVTopology | None = None + self.transfer_topo: TransferTopology | None = None - self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size} - self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size} # With heterogeneous TP, P must wait for all assigned D TP workers to # finish reading before safely freeing the blocks. self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int) @@ -355,7 +351,6 @@ class NixlConnectorWorker: self.block_size // kernel_block_size ) self.block_size = kernel_block_size - self._block_size[self.engine_id] = kernel_block_size self.num_blocks *= self._physical_blocks_per_logical_kv_block def _nixl_handshake( @@ -384,8 +379,8 @@ class NixlConnectorWorker: # Regardless, only handshake with the remote TP rank(s) that current # local rank will read from. Note that With homogeneous TP, # this happens to be the same single rank_i. - assert self.kv_topo is not None - p_remote_ranks = self.kv_topo.get_target_remote_ranks(remote_tp_size) + assert self.transfer_topo is not None + p_remote_ranks = self.transfer_topo.handshake_target_ranks(remote_tp_size) remote_rank_to_agent_name = {} path = make_zmq_path("tcp", host, port) @@ -649,11 +644,11 @@ class NixlConnectorWorker: def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in nixl.""" - self.kv_topo = TpKVTopology( + self.transfer_topo = TransferTopology( tp_rank=self.tp_rank, + tp_size=self.world_size, + block_size=self.block_size, engine_id=self.engine_id, - remote_tp_size=self._tp_size, # shared state - remote_block_size=self._block_size, # shared state is_mla=self.use_mla, total_num_kv_heads=self.model_config.get_total_num_kv_heads(), attn_backends=self.attn_backends, @@ -664,7 +659,7 @@ class NixlConnectorWorker: is_mamba=self._has_mamba, ) self.compat_hash = compute_nixl_compatibility_hash( - self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks + self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks ) if self.use_host_buffer: @@ -716,7 +711,7 @@ class NixlConnectorWorker: if isinstance(layer_spec, UniformTypeKVCacheSpecs): # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs layer_spec = layer_spec.kv_cache_specs[layer_name] - cache_list = self.kv_topo.get_transfer_cache_regions( + cache_list = self.transfer_topo.get_transfer_cache_regions( cache_or_caches, layer_spec ) # `layer_spec.page_size_bytes` only accounts for logical page_size, that is @@ -729,7 +724,7 @@ class NixlConnectorWorker: ) # For when registering multiple tensors eg K/V in separate regions. physical_page_size = physical_page_size // len(cache_list) - if self.kv_topo._cross_layers_blocks: + if self.transfer_topo._cross_layers_blocks: # When cross-layers blocks are used, multiply by number of layers physical_page_size = physical_page_size * len( self.kv_cache_config.kv_cache_tensors @@ -793,7 +788,7 @@ class NixlConnectorWorker: self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses self.num_regions = len(caches_data) - if self.kv_topo.is_kv_layout_blocks_first: + if self.transfer_topo.is_kv_layout_blocks_first: # NOTE (NickLucche) When FlashInfer is used, memory is registered # with joint KV for each block. This minimizes the overhead in # registerMem allowing faster descs queries. In order to be able to @@ -817,7 +812,7 @@ class NixlConnectorWorker: self.dst_num_blocks[self.engine_id] = self.num_blocks if self._has_mamba: - self._mamba_phys_ratio[self.engine_id] = ( + self._physical_blocks_per_logical[self.engine_id] = ( self._physical_blocks_per_logical_kv_block ) logger.info( @@ -876,11 +871,13 @@ class NixlConnectorWorker: conv_offsets = self._conv_decomp.local_conv_offsets conv_size, ssm_size = self._mamba_ssm_size num_blocks = self._logical_num_blocks * block_size_ratio - phys_ratio = self._physical_blocks_per_logical_kv_block + physical_per_logical = self._physical_blocks_per_logical_kv_block result: list[tuple[int, int, int]] = [] for i, base_addr in enumerate(base_addresses): - page_stride = self.block_len_per_layer[i] // block_size_ratio * phys_ratio + page_stride = ( + self.block_len_per_layer[i] // block_size_ratio * physical_per_logical + ) for off, sz in conv_offsets: for blk in range(num_blocks): result.append( @@ -900,14 +897,14 @@ class NixlConnectorWorker: def _build_fa_remote_for_mamba( self, nixl_agent_meta: NixlAgentMetadata, - transfer_cfg: HeteroTPTransferConfig, block_size_ratio: int, - kv_topo: TpKVTopology, + transfer_topo: TransferTopology, + remote_engine_id: EngineId, ) -> list[tuple[int, int, int]]: """Build remote FA descriptors for mamba models. - Uses transfer_cfg for GQA-aware FA divisor and head-based rank offset - instead of the standard uniform tp_ratio split. + Uses TransferTopology for GQA-aware FA divisor and head-based rank + offset instead of the standard uniform tp_ratio split. """ assert block_size_ratio == 1, ( "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " @@ -915,7 +912,9 @@ class NixlConnectorWorker: ) # TODO (ZhanqiuHu): unify with register_remote_blocks when Mamba-HMA # hetero-TP logic stabilizes. - tp_ratio = transfer_cfg.tp_ratio + mamba_info = transfer_topo.get_engine_info(remote_engine_id) + assert isinstance(mamba_info, MambaEngineTransferInfo) + tp_ratio = transfer_topo.tp_ratio(mamba_info.remote_tp_size) result: list[tuple[int, int, int]] = [] for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): local_block_len = self.get_backend_aware_kv_block_len( @@ -926,9 +925,11 @@ class NixlConnectorWorker: local_block_len = remote_kv_block_len if tp_ratio < 0 and not self.use_mla: - local_block_len = local_block_len // transfer_cfg.physical_fa_num_reads + local_block_len = local_block_len // mamba_info.remote_num_fa_reads - rank_offset = transfer_cfg.fa_rank_offset(remote_kv_block_len) + rank_offset = transfer_topo.fa_rank_offset( + remote_engine_id, remote_kv_block_len + ) num_blocks = nixl_agent_meta.num_blocks page_size = nixl_agent_meta.block_lens[i] @@ -937,12 +938,12 @@ class NixlConnectorWorker: addr = base_addr + block_offset + rank_offset result.append((addr, local_block_len, nixl_agent_meta.device_id)) - if kv_topo.is_kv_layout_blocks_first: + if transfer_topo.is_kv_layout_blocks_first: second_split = self.get_backend_aware_kv_block_len( layer_idx=i, first_split=False, mamba_view=False ) if tp_ratio < 0 and not self.use_mla: - second_split = second_split // transfer_cfg.physical_fa_num_reads + second_split = second_split // mamba_info.remote_num_fa_reads for block_id in range(num_blocks): block_offset = block_id * page_size addr = base_addr + block_offset + rank_offset @@ -981,15 +982,17 @@ class NixlConnectorWorker: conv_offsets = [(0, xb_p), (xb_p, bb_p), (xb_p + bb_p, bb_p)] ssm_read_size = nixl_agent_meta.ssm_sizes[1] - remote_ratio = self._mamba_phys_ratio[nixl_agent_meta.engine_id] - num_blocks = nixl_agent_meta.num_blocks // remote_ratio + remote_physical_per_logical = self._physical_blocks_per_logical[ + nixl_agent_meta.engine_id + ] + num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical device_id = nixl_agent_meta.device_id result: list[tuple[int, int, int]] = [] # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case # block lengths vary across layers (e.g. MLA). for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): - page_stride = nixl_agent_meta.block_lens[i] * remote_ratio + page_stride = nixl_agent_meta.block_lens[i] * remote_physical_per_logical for off, sz in conv_offsets: for blk in range(num_blocks): result.append((base_addr + blk * page_stride + off, sz, device_id)) @@ -1019,8 +1022,8 @@ class NixlConnectorWorker: register another local_xfer_handler using remote block len to ensure data copy correctness. """ - assert self.kv_topo is not None - kv_topo = self.kv_topo + assert self.transfer_topo is not None + transfer_topo = self.transfer_topo block_size_ratio = self.block_size // block_size blocks_data: list[tuple[int, int, int]] = [] @@ -1051,7 +1054,7 @@ class NixlConnectorWorker: # (addr, len, device id) blocks_data.append((addr, kv_block_len, self.device_id)) - if kv_topo.is_kv_layout_blocks_first: + if transfer_topo.is_kv_layout_blocks_first: second_split = self.get_backend_aware_kv_block_len( layer_idx=i, first_split=False, mamba_view=mamba ) @@ -1153,11 +1156,29 @@ class NixlConnectorWorker: ) return self._remote_agents[engine_id][remote_tp_rank] - ### Register remote agent metadata - if engine_id not in self._tp_size: - self._tp_size[engine_id] = remote_tp_size - if engine_id not in self._block_size: - self._block_size[engine_id] = nixl_agent_meta.block_size + ### Register remote engine in TransferTopology (idempotent). + assert self.transfer_topo is not None + transfer_topo = self.transfer_topo + physical_blocks_per_logical = ( + compute_physical_blocks_per_logical( + nixl_agent_meta.ssm_sizes, + nixl_agent_meta.block_lens[0], + ) + if self._has_mamba + else 1 + ) + transfer_topo.register_remote_engine( + remote_engine_id=engine_id, + remote_tp_size=remote_tp_size, + remote_block_size=nixl_agent_meta.block_size, + remote_block_len=nixl_agent_meta.block_lens[0], + remote_physical_blocks_per_logical=physical_blocks_per_logical, + local_block_len=self.block_len_per_layer[0], + ) + if self._has_mamba and engine_id not in self._physical_blocks_per_logical: + self._physical_blocks_per_logical[engine_id] = physical_blocks_per_logical + + logger.info("Transfer plan: %s", transfer_topo.describe(engine_id)) remote_agent_name = self.nixl_wrapper.add_remote_agent( nixl_agent_meta.agent_metadata @@ -1170,16 +1191,10 @@ class NixlConnectorWorker: # remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12| # local origin:| 0| 1| 8| 12| # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15| - assert self.kv_topo is not None - kv_topo = self.kv_topo - block_size_ratio = kv_topo.block_size_ratio_from_engine_id(engine_id) + block_size_ratio = transfer_topo.block_size_ratio(nixl_agent_meta.block_size) if engine_id not in self.dst_num_blocks: self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks - if self._has_mamba: - self._mamba_phys_ratio[engine_id] = compute_mamba_phys_ratio( - nixl_agent_meta.ssm_sizes, nixl_agent_meta.block_lens[0] - ) # Keep track of remote agent kv caches base addresses. self.kv_caches_base_addr[engine_id][remote_tp_rank] = ( @@ -1189,28 +1204,13 @@ class NixlConnectorWorker: # This is 1 when P and D `--tensor-parallel-size` match. Otherwise, # this is the ratio between the two sizes. - tp_ratio = self.kv_topo.tp_ratio_from_engine_id(engine_id) + tp_ratio = transfer_topo.tp_ratio(remote_tp_size) # Handle tp_size>num_kv_heads: replicate KV cache. indexes_into_remote = ( - not self.kv_topo.replicates_kv_cache(engine_id) and tp_ratio > 0 + not transfer_topo.replicates_kv_cache(engine_id) and tp_ratio > 0 ) - # Create transfer config (single source of truth for descriptor sizes). - if self._has_mamba and engine_id not in self._transfer_configs: - self._transfer_configs[engine_id] = HeteroTPTransferConfig( - tp_ratio=tp_ratio, - K=kv_topo.total_num_kv_heads, - d_tp=self.world_size, - p_tp=remote_tp_size, - d_rank=self.tp_rank, - use_mla=self.use_mla, - d_block_len=self.block_len_per_layer[0], - p_block_len=nixl_agent_meta.block_lens[0], - is_blocks_first=kv_topo.is_kv_layout_blocks_first, - ) - logger.info("Created %s", self._transfer_configs[engine_id].describe()) - logger.debug( "Registering remote agent (%s, rank %s) memory regions with tp_ratio %s", engine_id, @@ -1231,12 +1231,10 @@ class NixlConnectorWorker: self.src_xfer_handles_by_tp_ratio[tp_ratio] = [] if self._has_mamba: - transfer_cfg = self._transfer_configs.get(engine_id) - assert transfer_cfg is not None - if transfer_cfg.needs_split_handles: + if transfer_topo.needs_split_handles(engine_id): # Mamba-HMA: FA and Mamba use different split factors. - for handle_data in transfer_cfg.compute_split_handle_data( - self.src_blocks_data, self.num_descs, abs_tp + for handle_data in transfer_topo.compute_split_handle_data( + engine_id, self.src_blocks_data, self.num_descs, abs_tp ): descs = self.nixl_wrapper.get_xfer_descs( handle_data, self.nixl_memory_type @@ -1247,12 +1245,8 @@ class NixlConnectorWorker: self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle) logger.info( - "Mamba-HMA split handles: targets=%s, fa_reads=%s, " - "fa_entry=%s, mamba_reads=%s, num_descs=%s", - transfer_cfg.transfer_targets, - transfer_cfg.physical_fa_num_reads, - transfer_cfg.fa_entry_size, - transfer_cfg.mamba_num_reads, + "Mamba-HMA split handles: %s, num_descs=%s", + transfer_topo.describe(engine_id), self.num_descs, ) else: @@ -1321,7 +1315,7 @@ class NixlConnectorWorker: (addr, local_block_len, nixl_agent_meta.device_id) ) - if kv_topo.is_kv_layout_blocks_first: + if transfer_topo.is_kv_layout_blocks_first: # With FlashInfer index V separately to allow head splitting. second_split = self.get_backend_aware_kv_block_len( layer_idx=i, first_split=False, mamba_view=mamba @@ -1360,14 +1354,12 @@ class NixlConnectorWorker: engine_id, remote_tp_rank, ) - transfer_cfg = self._transfer_configs.get(engine_id) - assert transfer_cfg is not None blocks_data.extend( self._build_fa_remote_for_mamba( nixl_agent_meta, - transfer_cfg, block_size_ratio, - kv_topo, + transfer_topo, + engine_id, ) ) blocks_data.extend( @@ -1403,18 +1395,19 @@ class NixlConnectorWorker: """ remote_engine_id = nixl_agent_meta.engine_id - assert self._tp_size[remote_engine_id] == remote_tp_size - assert self.kv_topo is not None + assert self.transfer_topo is not None + remote_info = self.transfer_topo.get_engine_info(remote_engine_id) + assert remote_info.remote_tp_size == remote_tp_size - tp_ratio = self.kv_topo.tp_ratio_from_engine_id(remote_engine_id) - block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id( - remote_engine_id + tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size) + block_size_ratio = self.transfer_topo.block_size_ratio( + nixl_agent_meta.block_size ) # num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba. # Mamba models can have replicated FA KV with tp_ratio < 0. if not self._has_mamba: assert not ( - tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id) + tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id) ) if self._is_hma_required: @@ -1467,7 +1460,7 @@ class NixlConnectorWorker: if ( abs(tp_ratio) != 1 and not self.use_mla - and not self.kv_topo.is_kv_replicated(remote_engine_id) + and not self.transfer_topo.is_kv_replicated(remote_engine_id) and kv_cache_layout != "HND" and not self.enable_permute_local_kv ): @@ -1478,7 +1471,7 @@ class NixlConnectorWorker: # Block len can only vary across layers when using MLA. remote_block_len = nixl_agent_meta.block_lens[0] - if self.use_mla or self.kv_topo.is_kv_replicated(remote_engine_id): + if self.use_mla or self.transfer_topo.is_kv_replicated(remote_engine_id): # With replicated KV cache, only the number of blocks can differ. # TODO (ZhanqiuHu): For mamba models, validate FA and mamba # block_lens separately. @@ -1594,7 +1587,7 @@ class NixlConnectorWorker: if len(self.device_kv_caches) == 0: return assert block_size_ratio >= 1, "Only nP < nD supported currently." - assert self.kv_topo is not None + assert self.transfer_topo is not None if self.enable_permute_local_kv and block_size_ratio > 1: logger.debug( "Post-processing device kv cache on receive by converting " @@ -1614,7 +1607,7 @@ class NixlConnectorWorker: block_size_ratio, ) - split_k_and_v = self.kv_topo.split_k_and_v + split_k_and_v = self.transfer_topo.split_k_and_v for block_ids in block_ids_list: indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) @@ -1661,7 +1654,7 @@ class NixlConnectorWorker: The scheduler process (via the MultiprocExecutor) will use this output to track which workers are done. """ - assert self.kv_topo is not None + assert self.transfer_topo is not None done_sending = self._get_new_notifs() done_recving = self._pop_done_transfers(self._recving_transfers) @@ -1689,8 +1682,9 @@ class NixlConnectorWorker: self.sync_recved_kv_to_device(req_id, meta) # post processing for heteroblocksize - block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id( - meta.remote.engine_id + remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id) + block_size_ratio = self.transfer_topo.block_size_ratio( + remote_info.remote_block_size ) if not self.use_mla and ( block_size_ratio > 1 or self.enable_permute_local_kv @@ -1741,7 +1735,7 @@ class NixlConnectorWorker: are reading from the same producer (heterogeneous TP scenario), wait for all consumers to be done pulling. """ - assert self.kv_topo is not None + assert self.transfer_topo is not None notified_req_ids: set[str] = set() for notifs in self.nixl_wrapper.get_new_notifs().values(): for notif in notifs: @@ -1760,7 +1754,7 @@ class NixlConnectorWorker: # NOTE: `tp_ratio` is the opposite when swapping local<>remote n_consumers = int(tp_size) - tp_ratio = self.kv_topo.tp_ratio(n_consumers) + tp_ratio = self.transfer_topo.tp_ratio(n_consumers) # Number of reads *per producer* to wait for. # When remote D TP > local P TP we expect `tp_ratio` reads. @@ -1901,17 +1895,17 @@ class NixlConnectorWorker: self._reqs_to_send[req_id] = expiration_time def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): - assert meta.remote is not None and self.kv_topo is not None - remote_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id( - meta.remote.engine_id - ) - tp_ratio = self.kv_topo.tp_ratio_from_engine_id(meta.remote.engine_id) + assert meta.remote is not None and self.transfer_topo is not None + engine_id = meta.remote.engine_id + remote_ranks = self.transfer_topo.target_remote_ranks(engine_id) + remote_info = self.transfer_topo.get_engine_info(engine_id) + tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size) if self._has_mamba: # Expand remote logical → kernel block IDs. meta.remote.block_ids = self._logical_to_remote_kernel_block_ids( meta.remote.block_ids, - self._mamba_phys_ratio[meta.remote.engine_id], + self._physical_blocks_per_logical[meta.remote.engine_id], ) else: meta.remote.block_ids = self._logical_to_kernel_block_ids( @@ -1924,7 +1918,7 @@ class NixlConnectorWorker: # the first remote rank (cache is duplicated).. break - remote_block_size = self.kv_topo.remote_block_size[meta.remote.engine_id] + remote_block_size = remote_info.remote_block_size logger.debug( "Remote agent %s available, calling _read_blocks" " on remote rank %s with remote block size %s for req %s", @@ -1955,9 +1949,8 @@ class NixlConnectorWorker: remote_ids: BlockIds = meta.remote.block_ids if self._has_mamba: # Mamba-HMA: zero out FA groups for P ranks outside fa_read_targets. - transfer_cfg = self._transfer_configs.get(meta.remote.engine_id) - assert transfer_cfg is not None - local_ids, remote_ids = transfer_cfg.filter_block_ids_for_rank( + local_ids, remote_ids = self.transfer_topo.filter_block_ids_for_rank( + engine_id, remote_rank, local_ids, remote_ids, @@ -1999,8 +1992,11 @@ class NixlConnectorWorker: Post a READ point-to-point xfer request from a single local worker to a single remote worker. """ - assert self.kv_topo is not None - block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id) + assert self.transfer_topo is not None + remote_info = self.transfer_topo.get_engine_info(dst_engine_id) + block_size_ratio = self.transfer_topo.block_size_ratio( + remote_info.remote_block_size + ) if block_size_ratio > 1: # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups. assert not self._is_hma_required @@ -2190,8 +2186,8 @@ class NixlConnectorWorker: # This is like having two "low-level views" of the same storage. # `num_fa_descs` offset must be computed per-engine since P and D can # have different num_blocks (and thus different FA descs counts). - ratio = self._mamba_phys_ratio[engine_id] - logical_blocks = num_blocks // ratio + physical_per_logical = self._physical_blocks_per_logical[engine_id] + logical_blocks = num_blocks // physical_per_logical num_fa_descs = self.num_regions * num_blocks # 3-read mamba: 4 regions per unique cache tensor (x, B, C, ssm). mamba_region_ids = np.arange(len(self.block_len_per_layer) * 4)[:, None] @@ -2234,21 +2230,22 @@ class NixlConnectorWorker: ] def _logical_to_remote_kernel_block_ids( - self, block_ids: BlockIds, remote_ratio: int + self, block_ids: BlockIds, remote_physical_per_logical: int ) -> BlockIds: """Map logical block IDs to physical kernel block IDs on the remote. Args: block_ids: per-group lists of logical block IDs. - remote_ratio: remote engine's physical blocks per logical block. + remote_physical_per_logical: remote engine's physical blocks + per logical block. Returns: Same structure with FA groups expanded (each logical block L - becomes kernel blocks [L*remote_ratio .. L*remote_ratio + - local_ratio - 1]). Mamba groups are passed through unchanged. + becomes kernel blocks [L*ratio .. L*ratio + local_ratio - 1]). + Mamba groups are passed through unchanged. """ local_ratio = self._physical_blocks_per_logical_kv_block - if remote_ratio == 1: + if remote_physical_per_logical == 1: return block_ids local_arange = np.arange(local_ratio).reshape(1, -1) group_specs = self.kv_cache_config.kv_cache_groups @@ -2256,7 +2253,7 @@ class NixlConnectorWorker: for i, group in enumerate(block_ids): if not isinstance(group_specs[i].kv_cache_spec, MambaSpec): arr = np.array(group).reshape(-1, 1) - expanded = (arr * remote_ratio + local_arange).flatten() + expanded = (arr * remote_physical_per_logical + local_arange).flatten() result.append(expanded.tolist()) else: # Mamba blocks are 1:1 logical-to-physical (no expansion). @@ -2296,8 +2293,8 @@ class NixlConnectorWorker: +-------------------+ +--------------------+ |1st_split-2nd_split| |1st_split-2nd_split | """ - assert self.kv_topo is not None - if self.kv_topo.is_kv_layout_blocks_first: + assert self.transfer_topo is not None + if self.transfer_topo.is_kv_layout_blocks_first: # For indexing only half (either just the K or V part). if mamba_view: # NOTE (NickLucche) Mamba Opt: this is already skipping the padding so diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py index c8a5e10344b..309426814c6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py @@ -151,7 +151,9 @@ def derive_mamba_conv_split( ) -def compute_mamba_phys_ratio(ssm_sizes: tuple[int, ...], block_len: int) -> int: +def compute_physical_blocks_per_logical( + ssm_sizes: tuple[int, ...], block_len: int +) -> int: """Derive _physical_blocks_per_logical_kv_block from remote metadata. The remote engine's ratio is not sent directly in the handshake, so we From e06de7f0057fe1dfcc4fb039b52a61d39a079c4c Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Mon, 20 Apr 2026 20:57:11 +0800 Subject: [PATCH 140/150] [XPU] enable triton attention test on XPU by removing cuda device binding (#39627) Signed-off-by: Yan Ma --- .../attention/test_triton_decode_attention.py | 55 +++++++++++-------- .../test_triton_prefill_attention.py | 27 +++++---- .../test_triton_unified_attention.py | 6 +- 3 files changed, 52 insertions(+), 36 deletions(-) diff --git a/tests/kernels/attention/test_triton_decode_attention.py b/tests/kernels/attention/test_triton_decode_attention.py index a9b88162944..81e8bb17e7b 100644 --- a/tests/kernels/attention/test_triton_decode_attention.py +++ b/tests/kernels/attention/test_triton_decode_attention.py @@ -4,9 +4,12 @@ import pytest import torch +from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd +DEVICE_TYPE = current_platform.device_type + @pytest.mark.parametrize("B", [3, 5]) @pytest.mark.parametrize("L", [1027, 1025]) @@ -25,33 +28,35 @@ def test_decode_attention(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE): num_pages_per_batch = cdiv(seq_len, PAGE_SIZE) req_to_page = torch.randint( - 0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device="cuda" + 0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device=DEVICE_TYPE ) req_to_token = req_to_page * PAGE_SIZE req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE) - req_to_token = req_to_token + torch.arange(PAGE_SIZE, device="cuda").view(1, 1, -1) + req_to_token = req_to_token + torch.arange(PAGE_SIZE, device=DEVICE_TYPE).view( + 1, 1, -1 + ) req_to_token = req_to_token.view(B, -1) req_to_token = req_to_token[:, :seq_len].contiguous() # q represents the new token being generated, one per batch - q = torch.randn(B, H_Q, D_QK, dtype=dtype, device="cuda") + q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE) # k_buffer and v_buffer represent all previous tokens # Page size is 1. - k_buffer = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device="cuda") - v_buffer = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device="cuda") + k_buffer = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE) + v_buffer = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE) # o will have the same shape as q - o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda") + o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE) - lse = torch.zeros(B, H_Q, dtype=dtype, device="cuda") + lse = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE) - b_seq_len = torch.full((B,), seq_len, device="cuda") + b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE) attn_logits = torch.empty( (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, - device="cuda", + device=DEVICE_TYPE, ) # Call the original implementation. @@ -127,25 +132,27 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE) num_pages_per_batch = cdiv(seq_len, PAGE_SIZE) req_to_page = torch.randint( - 0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device="cuda" + 0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device=DEVICE_TYPE ) req_to_token = req_to_page * PAGE_SIZE req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE) - req_to_token = req_to_token + torch.arange(PAGE_SIZE, device="cuda").view(1, 1, -1) + req_to_token = req_to_token + torch.arange(PAGE_SIZE, device=DEVICE_TYPE).view( + 1, 1, -1 + ) req_to_token = req_to_token.view(B, -1) req_to_token = req_to_token[:, :seq_len].contiguous() - q = torch.randn(B, H_Q, D_QK, dtype=dtype, device="cuda") + q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE) # Create BF16 K/V as reference - k_bf16 = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device="cuda") - v_bf16 = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device="cuda") + k_bf16 = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE) + v_bf16 = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE) # --- BF16 reference --- - o_ref = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda") - lse_ref = torch.zeros(B, H_Q, dtype=dtype, device="cuda") + o_ref = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE) + lse_ref = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE) attn_logits = torch.empty( - (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device="cuda" + (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE ) if PAGE_SIZE == 1: @@ -156,7 +163,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE) o_ref, lse_ref, req_to_token, - b_seq_len=torch.full((B,), seq_len, device="cuda"), + b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE), attn_logits=attn_logits, num_kv_splits=num_kv_splits, sm_scale=sm_scale, @@ -171,7 +178,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE) o_ref, lse_ref, req_to_page, - b_seq_len=torch.full((B,), seq_len, device="cuda"), + b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE), attn_logits=attn_logits, num_kv_splits=num_kv_splits, sm_scale=sm_scale, @@ -182,10 +189,10 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE) k_fp8, k_scale = _quantize_to_fp8(k_bf16) v_fp8, v_scale = _quantize_to_fp8(v_bf16) - o_fp8 = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda") - lse_fp8 = torch.zeros(B, H_Q, dtype=dtype, device="cuda") + o_fp8 = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE) + lse_fp8 = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE) attn_logits_fp8 = torch.empty( - (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device="cuda" + (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE ) if PAGE_SIZE == 1: @@ -196,7 +203,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE) o_fp8, lse_fp8, req_to_token, - b_seq_len=torch.full((B,), seq_len, device="cuda"), + b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE), attn_logits=attn_logits_fp8, num_kv_splits=num_kv_splits, sm_scale=sm_scale, @@ -213,7 +220,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE) o_fp8, lse_fp8, req_to_page, - b_seq_len=torch.full((B,), seq_len, device="cuda"), + b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE), attn_logits=attn_logits_fp8, num_kv_splits=num_kv_splits, sm_scale=sm_scale, diff --git a/tests/kernels/attention/test_triton_prefill_attention.py b/tests/kernels/attention/test_triton_prefill_attention.py index f4505d91f5f..6316a926ae3 100644 --- a/tests/kernels/attention/test_triton_prefill_attention.py +++ b/tests/kernels/attention/test_triton_prefill_attention.py @@ -5,8 +5,11 @@ import pytest import torch import torch.nn.functional as F +from vllm.platforms import current_platform from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd +DEVICE_TYPE = current_platform.device_type + def ref_masked_attention( q: torch.Tensor, @@ -92,17 +95,19 @@ def test_context_attention( torch.manual_seed(42) # Generate random sequence lengths for each batch - seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda") + seq_lens = torch.randint( + max_seq_len // 2, max_seq_len + 1, (B,), device=DEVICE_TYPE + ) total_tokens = seq_lens.sum().item() # Create batch start locations - b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda") + b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE) b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0) # Create input tensors - q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device="cuda") - k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda") - v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda") + q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device=DEVICE_TYPE) + k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE) + v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE) o = torch.zeros_like(q) # Call Triton kernel @@ -169,17 +174,19 @@ def test_context_attention_sliding_window( torch.manual_seed(42) # Generate random sequence lengths for each batch - seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda") + seq_lens = torch.randint( + max_seq_len // 2, max_seq_len + 1, (B,), device=DEVICE_TYPE + ) total_tokens = seq_lens.sum().item() # Create batch start locations - b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda") + b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE) b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0) # Create input tensors - q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device="cuda") - k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda") - v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda") + q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device=DEVICE_TYPE) + k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE) + v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE) o = torch.zeros_like(q) # Call Triton kernel diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index 99cdc7ffa4a..ef52d80929c 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -10,6 +10,8 @@ from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.ops.triton_unified_attention import unified_attention +DEVICE_TYPE = current_platform.device_type + NUM_HEADS = [(4, 4), (8, 2), (5, 1)] HEAD_SIZES = [128, 256] BLOCK_SIZES = [16] @@ -114,7 +116,7 @@ def test_triton_unified_attn( q_dtype: torch.dtype | None, seq_threshold_3D: int, ) -> None: - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) set_random_seed(0) num_seqs = len(seq_lens) @@ -249,7 +251,7 @@ def test_triton_unified_attn_fp16_input_fp8_output( seq_threshold_3D: int, ) -> None: """Test with fp16 input and fp8 output using output_scale.""" - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) set_random_seed(0) num_seqs = len(seq_lens) From b82fc1364d313a222bde7e9ac897fd847ad7b05b Mon Sep 17 00:00:00 2001 From: aleksandaryanakiev Date: Mon, 20 Apr 2026 16:10:45 +0300 Subject: [PATCH 141/150] [Anthropic][Frontend] Added chat_template_kwargs to /v1/messages (#40125) Signed-off-by: Aleksandar Yanakiev Co-authored-by: Aleksandar Yanakiev --- vllm/entrypoints/anthropic/protocol.py | 16 ++++++++++++++++ vllm/entrypoints/anthropic/serving.py | 2 ++ 2 files changed, 18 insertions(+) diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py index 3445f709109..52eb77b5116 100644 --- a/vllm/entrypoints/anthropic/protocol.py +++ b/vllm/entrypoints/anthropic/protocol.py @@ -117,6 +117,13 @@ class AnthropicMessagesRequest(BaseModel): default=None, description="KVTransfer parameters used for disaggregated serving.", ) + chat_template_kwargs: dict[str, Any] | None = Field( + default=None, + description=( + "Additional keyword args to pass to the chat template renderer. " + "Will be accessible by the template." + ), + ) @field_validator("model") @classmethod @@ -212,6 +219,15 @@ class AnthropicCountTokensRequest(BaseModel): tool_choice: AnthropicToolChoice | None = None tools: list[AnthropicTool] | None = None + # vLLM-specific fields that are not in Anthropic spec + chat_template_kwargs: dict[str, Any] | None = Field( + default=None, + description=( + "Additional keyword args to pass to the chat template renderer. " + "Will be accessible by the template." + ), + ) + @field_validator("model") @classmethod def validate_model(cls, v): diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 5136e2b0fb0..939f5a7ed4c 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -323,6 +323,7 @@ class AnthropicServingMessages(OpenAIServingChat): return ChatCompletionRequest( model=anthropic_request.model, messages=openai_messages, + chat_template_kwargs=anthropic_request.chat_template_kwargs, ) return ChatCompletionRequest( @@ -335,6 +336,7 @@ class AnthropicServingMessages(OpenAIServingChat): top_p=anthropic_request.top_p, top_k=anthropic_request.top_k, kv_transfer_params=anthropic_request.kv_transfer_params, + chat_template_kwargs=anthropic_request.chat_template_kwargs, ) @classmethod From a023edfa5bc1982ba5c6365ea2e36c6dfc48a9ef Mon Sep 17 00:00:00 2001 From: Galigator Date: Mon, 20 Apr 2026 15:19:57 +0200 Subject: [PATCH 142/150] [bugfix] Use only onlines CPUs in lscpu (#40161) Signed-off-by: kse Co-authored-by: kse --- vllm/utils/cpu_resource_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/utils/cpu_resource_utils.py b/vllm/utils/cpu_resource_utils.py index bf3d84fb68d..4a56e7f6433 100644 --- a/vllm/utils/cpu_resource_utils.py +++ b/vllm/utils/cpu_resource_utils.py @@ -155,7 +155,7 @@ def _get_cpu_list() -> list[LogicalCPUInfo]: return [LogicalCPUInfo(i, i, 0) for i in range(cpu_count)] lscpu_output = subprocess.check_output( - "lscpu -J -e=CPU,CORE,NODE", shell=True, text=True + "lscpu --json --extended=CPU,CORE,NODE --online", shell=True, text=True ) # For platform without NUMA, replace '-' to '0' From 38fa87cacadc73aec9b28fea52fa70a31070cec4 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Mon, 20 Apr 2026 09:26:12 -0400 Subject: [PATCH 143/150] mxfp8 online quant move to new frontend (#40152) Signed-off-by: Vasiliy Kuznetsov --- docs/features/quantization/online.md | 7 +- vllm/config/quantization.py | 3 +- .../layers/quantization/__init__.py | 4 +- .../model_executor/layers/quantization/fp8.py | 8 - .../layers/quantization/online/base.py | 8 + .../layers/quantization/{ => online}/mxfp8.py | 202 ++++++++---------- .../model_executor/warmup/deep_gemm_warmup.py | 2 +- 7 files changed, 108 insertions(+), 126 deletions(-) rename vllm/model_executor/layers/quantization/{ => online}/mxfp8.py (57%) diff --git a/docs/features/quantization/online.md b/docs/features/quantization/online.md index aa02366d2d5..913da250ee6 100644 --- a/docs/features/quantization/online.md +++ b/docs/features/quantization/online.md @@ -17,6 +17,9 @@ llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_tensor") # Per-block FP8 quantization (128x128 block scaling for weights and 1x128 block scaling for activations) llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block") + +# MXFP8 quantization for weights and activations +llm = LLM("meta-llama/Llama-3.1-8B", quantization="mxfp8") ``` Or with the CLI: @@ -24,6 +27,7 @@ Or with the CLI: ```bash vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block +vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8 ``` ## Supported Schemes @@ -32,8 +36,7 @@ vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block | ------ | ------------- | ------------------ | ----- | | `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance | | `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | | - -Support for additional schemes will be added in future versions of vllm. +| `mxfp8` | fp8_e4m3 data, e8m0 per-1x32-block scale | fp8_e4m3 data, e8m0 per-1x32-block scale | Requires SM 100+ (Blackwell or newer) for w8a8, other GPUs use a w8a16 fallback | ## Advanced Configuration diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index 02df08b8b48..05e91f94025 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -23,7 +23,8 @@ class OnlineQuantScheme(Enum): # Linear layers remain unquantized. INT8_PER_CHANNEL_WEIGHT_ONLY = "int8_per_channel_weight_only" - # TODO(future PRs): add more online quant schemes here: mxfp8, etc + # mxfp8, weights scaled in blocks of 1x32 elements (microscaling FP8) + MXFP8 = "mxfp8" @config diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index 352d04ea208..aec9f7c3d71 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -31,7 +31,6 @@ QuantizationMethods = Literal[ "inc", "mxfp4", "gpt_oss_mxfp4", - "mxfp8", "cpu_awq", "online", # Below are values of the OnlineQuantScheme enum, specified as strings to @@ -41,6 +40,7 @@ QuantizationMethods = Literal[ "fp8_per_tensor", "fp8_per_block", "int8_per_channel_weight_only", + "mxfp8", ] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) @@ -135,7 +135,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: ) from .moe_wna16 import MoeWNA16Config from .mxfp4 import GptOssMxfp4Config, Mxfp4Config - from .mxfp8 import Mxfp8Config from .online.base import OnlineQuantizationConfig from .torchao import TorchAOConfig @@ -162,7 +161,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "inc": INCConfig, "mxfp4": Mxfp4Config, "gpt_oss_mxfp4": GptOssMxfp4Config, - "mxfp8": Mxfp8Config, "cpu_awq": CPUAWQConfig, "online": OnlineQuantizationConfig, } diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 94ed4f7c97b..4473bb8b2cc 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -515,14 +515,6 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod): initialize_online_processing(layer) - # TODO: remove this check once the following RFC is resolved. - # https://github.com/vllm-project/vllm/issues/33314 - # Subclasses (e.g. Mxfp8OnlineLinearMethod) only need the weight - # registration above and manage their own kernel, so skip fp8_linear - # kernel creation for them. - if type(self) is not Fp8OnlineLinearMethod: - return - self.fp8_linear = init_fp8_linear_kernel( activation_quant_key=self.activation_quant_key, weight_quant_key=self.weight_quant_key, diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index 1f923c07ff1..315dcfacffc 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -37,6 +37,10 @@ from vllm.model_executor.layers.quantization.online.fp8 import ( from vllm.model_executor.layers.quantization.online.int8 import ( Int8OnlineMoEMethod, ) +from vllm.model_executor.layers.quantization.online.mxfp8 import ( + Mxfp8OnlineLinearMethod, + Mxfp8OnlineMoEMethod, +) logger = init_logger(__name__) @@ -110,6 +114,8 @@ class OnlineQuantizationConfig(QuantizationConfig): return UnquantizedLinearMethod() elif linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK: return Fp8PerBlockOnlineLinearMethod() + elif linear_scheme == OnlineQuantScheme.MXFP8: + return Mxfp8OnlineLinearMethod() else: return Fp8PerTensorOnlineLinearMethod() elif isinstance(layer, FusedMoE): @@ -125,6 +131,8 @@ class OnlineQuantizationConfig(QuantizationConfig): return Int8OnlineMoEMethod(layer=layer) elif moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK: return Fp8PerBlockOnlineMoEMethod(layer=layer) + elif moe_scheme == OnlineQuantScheme.MXFP8: + return Mxfp8OnlineMoEMethod(layer=layer) else: return Fp8PerTensorOnlineMoEMethod(layer=layer) return None diff --git a/vllm/model_executor/layers/quantization/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py similarity index 57% rename from vllm/model_executor/layers/quantization/mxfp8.py rename to vllm/model_executor/layers/quantization/online/mxfp8.py index 5acf843f108..39a32604442 100644 --- a/vllm/model_executor/layers/quantization/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -1,131 +1,47 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Online MXFP8 (microscaling FP8, block-32) quantization config and methods.""" +"""Online MXFP8 (microscaling FP8, block-32) quantization methods.""" -from typing import Any +from typing import TYPE_CHECKING import torch from torch.nn import Module -from vllm.logger import init_logger +if TYPE_CHECKING: + import vllm.model_executor.layers.fused_moe.modular_kernel as mk + from vllm.model_executor.layers.fused_moe import FusedMoE + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, + ) + from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend + from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, - FusedMoEMethodBase, -) -from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( select_mxfp8_moe_backend, ) -from vllm.model_executor.layers.linear import ( - LinearBase, - UnquantizedLinearMethod, +from vllm.model_executor.layers.quantization.online.fp8 import ( + _Fp8OnlineLinearBase, ) -from vllm.model_executor.layers.quantization import QuantizationMethods -from vllm.model_executor.layers.quantization.base_config import ( - QuantizeMethodBase, -) -from vllm.model_executor.layers.quantization.fp8 import ( - Fp8Config, - Fp8KVCacheMethod, - Fp8OnlineLinearMethod, - Fp8OnlineMoEMethod, +from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, ) from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_BLOCK_SIZE, mxfp8_e4m3_quantize, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - is_layer_skipped, -) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform -logger = init_logger(__name__) - -class Mxfp8Config(Fp8Config): - """Config class for online MXFP8 MoE quantization.""" - - def __init__( - self, - activation_scheme: str = "dynamic", - ignored_layers: list[str] | None = None, - ) -> None: - if activation_scheme != "dynamic": - raise ValueError("mxfp8 only supports dynamic activation scheme.") - super().__init__( - is_checkpoint_fp8_serialized=False, - activation_scheme=activation_scheme, - ignored_layers=ignored_layers, - weight_block_size=None, - ) - - @classmethod - def get_name(cls) -> QuantizationMethods: - return "mxfp8" - - @classmethod - def get_min_capability(cls) -> int: - # Marlin kernel supports MXFP8 on SM80+ - return 80 - - @classmethod - def from_config(cls, config: dict[str, Any]) -> "Mxfp8Config": - activation_scheme = cls.get_from_keys_or( - config, ["activation_scheme"], "dynamic" - ) - ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) - if not ignored_layers: - ignored_layers = cls.get_from_keys_or( - config, ["modules_to_not_convert"], None - ) - return cls( - activation_scheme=activation_scheme, - ignored_layers=ignored_layers, - ) - - def get_quant_method( - self, layer: torch.nn.Module, prefix: str - ) -> "QuantizeMethodBase | None": - if isinstance(layer, LinearBase): - if is_layer_skipped( - prefix=prefix, - ignored_layers=self.ignored_layers, - fused_mapping=self.packed_modules_mapping, - skip_with_substr=True, - ): - return UnquantizedLinearMethod() - return Mxfp8OnlineLinearMethod(self) - elif isinstance(layer, FusedMoE): - if is_layer_skipped( - prefix=prefix, - ignored_layers=self.ignored_layers, - fused_mapping=self.packed_modules_mapping, - skip_with_substr=True, - ): - return UnquantizedFusedMoEMethod(layer.moe_config) - return Mxfp8OnlineMoEMethod(self, layer) - elif isinstance(layer, Attention): - return Fp8KVCacheMethod(self) - return None - - -class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod): +class Mxfp8OnlineLinearMethod(_Fp8OnlineLinearBase): """Online MXFP8 linear method. Loads bf16/fp16 checkpoints and quantizes weights to MXFP8 (microscaling FP8 with block-32 scales) during weight loading. - - Args: - quant_config: The MXFP8 quantization config. """ - uses_meta_device: bool = True - - def __init__(self, quant_config: "Mxfp8Config"): - self.quant_config = quant_config + def __init__(self): + super().__init__() self.kernel = init_mxfp8_linear_kernel() def create_weights( @@ -178,19 +94,15 @@ class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod): return self.kernel.apply_weights(layer, x, bias) -class Mxfp8OnlineMoEMethod(Fp8OnlineMoEMethod): +class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase): """MoE method for online MXFP8 (block) quantization.""" - uses_meta_device: bool = True + fp8_backend: "Fp8MoeBackend" + experts_cls: "type[mk.FusedMoEExperts] | None" - def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module): - FusedMoEMethodBase.__init__(self, layer.moe_config) - self.quant_config = quant_config - assert not quant_config.is_checkpoint_fp8_serialized - assert quant_config.activation_scheme == "dynamic" - - self.weight_block_size = [1, MXFP8_BLOCK_SIZE] - self.block_quant = True + def __init__(self, *, layer: torch.nn.Module): + super().__init__(layer.moe_config) + self.weight_block_size: list[int] = [1, MXFP8_BLOCK_SIZE] self.weight_scale_name = "weight_scale" self.fp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe) @@ -247,6 +159,74 @@ class Mxfp8OnlineMoEMethod(Fp8OnlineMoEMethod): return w_quant, w_scales + def _setup_kernel( + self, + layer: "FusedMoE", + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_input_scale: torch.Tensor | None, + w2_input_scale: torch.Tensor | None, + ) -> None: + from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + convert_to_fp8_moe_kernel_format, + make_fp8_moe_kernel, + ) + + # Shuffle weights to runtime format. + w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( + fp8_backend=self.fp8_backend, + layer=layer, + w13=w13, + w2=w2, + w13_scale=w13_scale, + w2_scale=w2_scale, + w13_input_scale=w13_input_scale, + w2_input_scale=w2_input_scale, + ) + + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale) + replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale) + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.moe_quant_config: + assert self.experts_cls is not None + self.moe_kernel = make_fp8_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + fp8_backend=self.fp8_backend, + experts_cls=self.experts_cls, + routing_tables=layer._maybe_init_expert_routing_tables(), + shared_experts=layer.shared_experts, + ) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> "FusedMoEQuantConfig": + from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + make_fp8_moe_quant_config, + ) + + w1_scale = getattr(layer, f"w13_{self.weight_scale_name}") + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") + a1_scale = layer.w13_input_scale + a2_scale = layer.w2_input_scale + + quant_config = make_fp8_moe_quant_config( + fp8_backend=self.fp8_backend, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=self.weight_block_size, + ) + + self._maybe_inject_biases(quant_config, layer) + return quant_config + def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index e41b2c5e100..a352cc116f5 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -19,7 +19,7 @@ from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import ( ) from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod -from vllm.model_executor.layers.quantization.mxfp8 import Mxfp8OnlineLinearMethod +from vllm.model_executor.layers.quantization.online.mxfp8 import Mxfp8OnlineLinearMethod from vllm.tracing import instrument from vllm.utils.deep_gemm import ( fp8_gemm_nt, From def8f52200151c801867dde9ce27f829bce00105 Mon Sep 17 00:00:00 2001 From: Sage Moore Date: Mon, 20 Apr 2026 07:22:54 -0700 Subject: [PATCH 144/150] [CI][EPLB] Add Async EPLB end-to-end integration test to CI (#40168) Signed-off-by: Sage Moore --- .../qwen30b_a3b_fp8_dp4_async_eplb.sh | 55 +++++++++++++++++++ .buildkite/test_areas/e2e_integration.yaml | 9 +++ 2 files changed, 64 insertions(+) create mode 100755 .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh diff --git a/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh b/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh new file mode 100755 index 00000000000..06743f16b68 --- /dev/null +++ b/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +set -euxo pipefail + +# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT] +THRESHOLD=${1:-0.8} +NUM_Q=${2:-1319} +PORT=${3:-8050} +OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled} +mkdir -p "${OUT_DIR}" + +wait_for_server() { + local port=$1 + timeout 600 bash -c ' + until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do + sleep 1 + done' +} + +MODEL="Qwen/Qwen3-30B-A3B-FP8" +BACK="allgather_reducescatter" + +cleanup() { + if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then + kill "${SERVER_PID}" 2>/dev/null || true + for _ in {1..20}; do + kill -0 "${SERVER_PID}" 2>/dev/null || break + sleep 0.5 + done + kill -9 "${SERVER_PID}" 2>/dev/null || true + fi +} +trap cleanup EXIT + +VLLM_DEEP_GEMM_WARMUP=skip \ +vllm serve "$MODEL" \ +--enforce-eager \ +--data-parallel-size 4 \ +--enable-expert-parallel \ +--enable-eplb \ +--all2all-backend "$BACK" \ +--eplb-config '{"window_size":20, "step_interval":100, "use_async":true}' \ +--trust-remote-code \ +--max-model-len 2048 \ +--port "$PORT" & +SERVER_PID=$! +wait_for_server "$PORT" + +TAG=$(echo "$MODEL" | tr '/: \\n' '_____') +OUT="${OUT_DIR}/${TAG}_${BACK}.json" +python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}" +python3 - <= ${THRESHOLD}, f"${MODEL} ${BACK} accuracy {acc}" +PY diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index 5b7f96bc7a2..857fefd268a 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -29,6 +29,15 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 +- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy + timeout_in_minutes: 60 + device: h100 + optional: true + num_devices: 4 + working_dir: "/vllm-workspace" + commands: + - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050 + - label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100) timeout_in_minutes: 60 device: h100 From 7243e02aa1c6e01ab1d2a50aa7e27ebc9b2c6719 Mon Sep 17 00:00:00 2001 From: larryli2-amd Date: Mon, 20 Apr 2026 22:44:43 +0800 Subject: [PATCH 145/150] [ROCm][Feature] Enable AITER MLA attention backend to work with Eagle3 speculative decoding on ROCm (#39616) Signed-off-by: larryli2-amd Co-authored-by: TJian --- docs/design/attention_backends.md | 2 +- .../attention/backends/mla/rocm_aiter_mla.py | 184 ++++++++++++------ 2 files changed, 129 insertions(+), 57 deletions(-) diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d10c23038a6..13c56a526fc 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -215,7 +215,7 @@ configuration. | `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 | 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` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | 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 | diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 351fdba085a..d1a38322ae2 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -44,7 +44,11 @@ class AiterMLABackend(MLACommonBackend): @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [1] + # The aiter MLA decode kernel always operates with page_size=1 + # internally (the wrapper flattens kv_buffer via .view(-1, 1, 1, H)). + # We support any kernel_block_size by expanding block-level indices + # into per-token flat indices in the metadata builder. + return [MultipleOf(1)] @staticmethod def get_name() -> str: @@ -74,6 +78,8 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata): attn_out_dtype: torch.dtype = torch.bfloat16 # The max query output length: int max_qo_len: int | None = None + # Whether persistent MLA metadata was computed (only for qseqlen=1) + has_persistent_metadata: bool = False @dataclass @@ -105,7 +111,16 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): self.compilation_config = vllm_config.compilation_config self.decode_attn_out_dtype = vllm_config.model_config.dtype - # kernel block size is always 1. + + # Store the kernel block size from the spec. When kernel_block_size=1 + # (no spec-dec), behavior is identical to the original. When > 1 + # (e.g. 16 with Eagle3), we expand block-level indices into per-token + # flat indices since the aiter kernel always uses page_size=1 internally. + self.kernel_block_size = kv_cache_spec.block_size + + # In the flat view (.view(-1,1,1,H)), each token is its own page, + # so max_num_pages_per_req = max_model_len regardless of + # kernel_block_size. max_num_pages_per_req = vllm_config.model_config.max_model_len max_num_reqs = vllm_config.scheduler_config.max_num_seqs max_num_pages = max_num_reqs * max_num_pages_per_req @@ -115,8 +130,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # so we can only use the persistent buffer if a cudagraph is actually # being used. - # paged_kv_last_page_len is always 1s (kernel block size is always 1), - # so we create it once and reuse slices in both eager and cudagraph modes. + # paged_kv_last_page_len is always 1s (the aiter kernel always sees + # page_size=1 after .view(-1,1,1,H) flattening), so we create it + # once and reuse slices in both eager and cudagraph modes. self.paged_kv_last_page_len = torch.ones( max_num_reqs, dtype=torch.int32, device=device ) @@ -196,14 +212,14 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): num_decode_tokens: int, dcp_tot_seq_lens_device: torch.Tensor | None, ) -> AiterMLADecodeMetadata: - # kernel block size is always 1, although the kv block size is not 1. device = self.device num_reqs = seq_lens_device.size(0) - # kernel block size is always 1, so each page has exactly 1 token. - # last_page_len is always 1 - just slice the pre-initialized buffer. + # The aiter kernel always operates with page_size=1 (the wrapper + # flattens kv_buffer). last_page_len is always 1. paged_kv_last_page_len = self.paged_kv_last_page_len[:num_reqs] + # indptr: cumsum of seq_lens (one page per token in the flat view) paged_kv_indptr = torch.cat( [ torch.zeros(1, dtype=seq_lens_device.dtype, device=device), @@ -215,11 +231,19 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): if self.compilation_config.cudagraph_mode.has_full_cudagraphs(): self.paged_kv_indices.fill_(-1) - _copy_page_indices_kernel[(num_reqs,)]( + + # Expand block_table entries into per-token flat indices. + # When kernel_block_size=1, this degrades to a direct copy (identical + # to the original _copy_page_indices_kernel). + # When kernel_block_size=K>1, block_table entry b covering K tokens + # gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1). + _expand_page_indices_kernel[(num_reqs,)]( self.paged_kv_indices, block_table_tensor, block_table_tensor.stride(0), paged_kv_indptr, + seq_lens_device, + KERNEL_BLOCK_SIZE=self.kernel_block_size, BLOCK_SIZE=1024, ) paged_kv_indices = self.paged_kv_indices @@ -245,27 +269,37 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): 0, num_reqs + 1, step=1, dtype=torch.int32, device=device ) - from aiter import get_mla_metadata_v1 + # The aiter MLA ASM kernel only supports qseqlen=1 (single-token + # decode). With speculative decoding, the verification step has + # qseqlen > 1 (e.g. 8 for spec7). get_mla_metadata_v1 calls + # get_heuristic_kernel_mla which fails for qseqlen > 1. + # We track whether persistent metadata was successfully computed + # so forward_mqa can skip passing it (falling back to the kernel + # computing its own metadata internally, like v0.18.0). + has_persistent_metadata = False + if max_qo_len == 1: + from aiter import get_mla_metadata_v1 - get_mla_metadata_v1( - qo_indptr, - paged_kv_indptr, - paged_kv_last_page_len, - self._num_attention_heads, - 1, - True, - self._mla_work_meta_data, - self._mla_work_info_set, - self._mla_work_indptr, - self._mla_reduce_indptr, - self._mla_reduce_final_map, - self._mla_reduce_partial_map, - page_size=1, - kv_granularity=16, - max_seqlen_qo=max_qo_len, - uni_seqlen_qo=max_qo_len, - fast_mode=True, - ) + get_mla_metadata_v1( + qo_indptr, + paged_kv_indptr, + paged_kv_last_page_len, + self._num_attention_heads, + 1, + True, + self._mla_work_meta_data, + self._mla_work_info_set, + self._mla_work_indptr, + self._mla_reduce_indptr, + self._mla_reduce_final_map, + self._mla_reduce_partial_map, + page_size=1, + kv_granularity=16, + max_seqlen_qo=max_qo_len, + uni_seqlen_qo=max_qo_len, + fast_mode=True, + ) + has_persistent_metadata = True attn_metadata = AiterMLADecodeMetadata( block_table=block_table_tensor, @@ -277,6 +311,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): dcp_tot_seq_lens=dcp_tot_seq_lens_device, max_qo_len=max_qo_len, attn_out_dtype=self.decode_attn_out_dtype, + has_persistent_metadata=has_persistent_metadata, ) return attn_metadata @@ -290,41 +325,67 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): attn_metadata = super().build( common_prefix_len, common_attn_metadata, fast_build ) - attn_metadata.work_meta_data = self._mla_work_meta_data - attn_metadata.work_indptr = self._mla_work_indptr - attn_metadata.work_info_set = self._mla_work_info_set - attn_metadata.reduce_indptr = self._mla_reduce_indptr - attn_metadata.reduce_final_map = self._mla_reduce_final_map - attn_metadata.reduce_partial_map = self._mla_reduce_partial_map + if ( + attn_metadata.decode is not None + and attn_metadata.decode.has_persistent_metadata + ): + attn_metadata.work_meta_data = self._mla_work_meta_data + attn_metadata.work_indptr = self._mla_work_indptr + attn_metadata.work_info_set = self._mla_work_info_set + attn_metadata.reduce_indptr = self._mla_reduce_indptr + attn_metadata.reduce_final_map = self._mla_reduce_final_map + attn_metadata.reduce_partial_map = self._mla_reduce_partial_map return attn_metadata @triton.jit -def _copy_page_indices_kernel( +def _expand_page_indices_kernel( page_indices, block_table, block_table_stride, - cu_num_blocks, + cu_num_tokens, + seq_lens, + KERNEL_BLOCK_SIZE: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): - """Copy block table rows into a flat page_indices buffer using indptr. - Avoids blocking boolean mask indexing (tensor[mask]) which has - data-dependent output size and forces sync. - This is the same kernel as introduced in backends/flashinfer.py. + """Expand block table entries into per-token flat page indices. + + The aiter MLA kernel always operates with page_size=1 internally + (kv_buffer is flattened via .view(-1, 1, 1, H)). This kernel converts + block-level indices from the block table into individual token positions + in the flattened KV buffer. + + When KERNEL_BLOCK_SIZE=1: block_idx=t, offset=0, flat=block_id + (equivalent to a direct copy -- no regression from the original kernel). + + When KERNEL_BLOCK_SIZE=K: block table entry b (covering K tokens) + is expanded to flat indices b*K, b*K+1, ..., b*K+(K-1). """ req_idx = tl.program_id(0) row_ptr = block_table + req_idx * block_table_stride - start_idx = tl.load(cu_num_blocks + req_idx) - end_idx = tl.load(cu_num_blocks + req_idx + 1) - num_blocks = end_idx - start_idx + start_idx = tl.load(cu_num_tokens + req_idx) + num_tokens = tl.load(seq_lens + req_idx) offset = tl.arange(0, BLOCK_SIZE) - for i in tl.range(0, num_blocks, BLOCK_SIZE): - block_ids = tl.load(row_ptr + i + offset, mask=i + offset < num_blocks) + for i in tl.range(0, num_tokens, BLOCK_SIZE): + token_offsets = i + offset + mask = token_offsets < num_tokens + + # Which block in the block table does this token belong to? + block_idx = token_offsets // KERNEL_BLOCK_SIZE + # Offset within that block + offset_in_block = token_offsets % KERNEL_BLOCK_SIZE + + # Load the block ID from the block table + block_ids = tl.load(row_ptr + block_idx, mask=mask) + + # Compute flat index in the flattened kv_buffer + flat_indices = block_ids * KERNEL_BLOCK_SIZE + offset_in_block + tl.store( - page_indices + start_idx + i + offset, - block_ids, - mask=i + offset < num_blocks, + page_indices + start_idx + token_offsets, + flat_indices, + mask=mask, ) @@ -426,6 +487,24 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2) + # Build kwargs for mla_decode_fwd. Pass persistent metadata only + # when it was successfully computed (qseqlen=1 decode steps). + # For multi-token verification steps (spec-dec), the kernel falls + # back to computing metadata internally. + mla_kwargs = dict( + q_scale=layer._q_scale, + kv_scale=layer._k_scale, + ) + if attn_metadata.work_meta_data is not None: + mla_kwargs.update( + work_meta_data=attn_metadata.work_meta_data, + work_indptr=attn_metadata.work_indptr, + work_info_set=attn_metadata.work_info_set, + reduce_indptr=attn_metadata.reduce_indptr, + reduce_final_map=attn_metadata.reduce_final_map, + reduce_partial_map=attn_metadata.reduce_partial_map, + ) + rocm_aiter_ops.mla_decode_fwd( q, kv_buffer, @@ -436,14 +515,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): attn_metadata.decode.paged_kv_indptr, attn_metadata.decode.paged_kv_indices, attn_metadata.decode.paged_kv_last_page_len, - q_scale=layer._q_scale, - kv_scale=layer._k_scale, - work_meta_data=attn_metadata.work_meta_data, - work_indptr=attn_metadata.work_indptr, - work_info_set=attn_metadata.work_info_set, - reduce_indptr=attn_metadata.reduce_indptr, - reduce_final_map=attn_metadata.reduce_final_map, - reduce_partial_map=attn_metadata.reduce_partial_map, + **mla_kwargs, ) if self._needs_head_repeat: From b42e878ec0182aafbaed5cc28708e99287ef6856 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 20 Apr 2026 10:52:32 -0400 Subject: [PATCH 146/150] [Bug] Fix dcp error message (#40053) Signed-off-by: yewentao256 --- vllm/v1/worker/cp_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index c5febda691d..05cca52fc0d 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -33,7 +33,7 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: "implementations to return the softmax LSE during decode, " f"but {layer_impl.__class__.__name__} does not. " "Try a different backend by setting " - "VLLM_ATTENTION_BACKEND or disable DCP." + "--attention-backend or disable DCP." ) if pcp_size > 1: From fb5635d3f90635ce9d1acbc975baab6e911a262b Mon Sep 17 00:00:00 2001 From: Rita Brugarolas Date: Mon, 20 Apr 2026 07:56:27 -0700 Subject: [PATCH 147/150] [ROCm] Add MLA dual RMS norm fusion (Q, KV) pass for DeepSeek/Kimi-K2 (#39242) Signed-off-by: Rita Brugarolas Brufau --- docs/design/fusions.md | 40 +++++ docs/design/optimization_levels.md | 1 + .../passes/test_fuse_mla_dual_rms_norm.py | 148 ++++++++++++++++++ vllm/_aiter_ops.py | 42 +++++ .../passes/fusion/rocm_aiter_fusion.py | 107 ++++++++++++- vllm/compilation/passes/pass_manager.py | 4 + vllm/config/compilation.py | 9 ++ vllm/config/vllm.py | 11 ++ 8 files changed, 361 insertions(+), 1 deletion(-) create mode 100644 tests/compile/passes/test_fuse_mla_dual_rms_norm.py diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 1df0eb6f439..afc7f351500 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -31,6 +31,7 @@ or just on the low or high end. | [RMSNorm + Quant](#rmsnorm--quantization-fuse_norm_quant) | `fuse_norm_quant` | RMSNorm (+residual add) → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always | | [SiLU+Mul + Quant](#silumul--quantization-fuse_act_quant) | `fuse_act_quant` | SiLU+Mul activation → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always | | [RMSNorm + Padding](#rmsnorm--padding-fuse_act_padding) | `fuse_act_padding` | Residual add + RMSNorm → padding | O1 (ROCm/AITER only) | TBD | No | Always | +| [MLA Dual RMSNorm](#mla-dual-rmsnorm-fuse_mla_dual_rms_norm) | `fuse_mla_dual_rms_norm` | Paired Q + KV RMSNorm → single kernel | O1 (ROCm/AITER only) | ~2% | No | Always | ## Support Matrix @@ -51,6 +52,7 @@ The table below lists the quantization schemes supported by each fusion on each | `fuse_norm_quant` | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | — | FP8 static, FP8 per-token, FP8 per-group | | `fuse_act_quant` | FP8 static, NVFP4 | FP8 static, FP8 per-group (128/64) | FP8 static, FP8 per-group (128/64) | — | FP8 per-group | | `fuse_act_padding` | — | — | — | — | FP16/BF16 | +| `fuse_mla_dual_rms_norm` | — | — | — | — | BF16 | \* `fuse_attn_quant` support depends on the attention backend in use; not all backends support fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant) @@ -381,6 +383,44 @@ when the hidden size is 2880 and AITER Triton GEMMs *not* enabled. - Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`RocmAiterTritonAddRMSNormPadFusionPass`) +### MLA Dual RMSNorm (`fuse_mla_dual_rms_norm`) + +!!! info + ROCm/AITER-only. Targeted at DeepSeek-V3 / Kimi-K2 MLA attention. + +!!! note + When the native implementation of `rms_norm` is used (the default on CUDA and + ROCm for now), Inductor's built-in fusion already handles merging these norms + automatically. This explicit pass targets the case where AITER's custom + `rms_norm` op is active, which Inductor cannot fuse on its own. + +**What it fuses.** Fuses the paired `q_a_layernorm` and `kv_a_layernorm` RMS norm +operations in MLA attention into a single `fused_qk_rmsnorm` HIP kernel call via AITER, +reducing kernel launch overhead from 2 launches to 1 per MLA layer. + +```text +# Unfused: +q_c, kv_lora = split(projected, [q_dim, kv_dim]) +kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim]) +q_c = rms_norm(q_c, q_weight, eps) +kv_c = rms_norm(kv_c, kv_weight, eps) + +# Fused: +q_c, kv_lora = split(projected, [q_dim, kv_dim]) +kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim]) +q_normed, kv_normed = fused_mla_dual_rms_norm( + q_c, q_weight, kv_c, kv_weight, eps1, eps2) +``` + +Requires: AMD ROCm with AITER enabled. Enabled by default at optimization level O1 and above +when AITER is available. + +**Code locations.** + +- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`) +- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`) +- AITER kernel: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442) + ## See Also - [Optimization Levels](optimization_levels.md) — high-level presets that set diff --git a/docs/design/optimization_levels.md b/docs/design/optimization_levels.md index 591978b542e..dd0936ca9e5 100644 --- a/docs/design/optimization_levels.md +++ b/docs/design/optimization_levels.md @@ -56,6 +56,7 @@ Fusions: - `-cc.pass_config.fuse_norm_quant=True`* - `-cc.pass_config.fuse_act_quant=True`* - `-cc.pass_config.fuse_act_padding=True`† +- `-cc.pass_config.fuse_mla_dual_rms_norm=True`† \* These fusions are only enabled when either op is using a custom kernel, otherwise Inductor fusion is better.
† These fusions are ROCm-only and require AITER. diff --git a/tests/compile/passes/test_fuse_mla_dual_rms_norm.py b/tests/compile/passes/test_fuse_mla_dual_rms_norm.py new file mode 100644 index 00000000000..080417c9896 --- /dev/null +++ b/tests/compile/passes/test_fuse_mla_dual_rms_norm.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Unit test for the MLADualRMSNormFusionPass. + +The pass fuses paired q/kv RMS norms in MLA attention into a single +fused_mla_dual_rms_norm op backed by AITER's fused_qk_rmsnorm kernel. +""" + +import pytest +import torch + +import vllm.config +from tests.compile.backend import TestBackend +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops +from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass +from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass +from vllm.config import ( + CompilationConfig, + CompilationMode, + ModelConfig, + PassConfig, + VllmConfig, +) +from vllm.model_executor.layers.layernorm import RMSNorm + +# MLA attention geometry for DeepSeek-V3 / Kimi-K2 +Q_DIM = 1536 +KV_C_DIM = 512 +K_PE_DIM = 64 +EPS = 1e-6 + + +class MLADualRMSNormTestModel(torch.nn.Module): + """ + Minimal model reproducing the MLA dual RMS norm pattern: + linear -> split([q_dim, kv_dim]) + +-- q_c (getitem 0) -> rms_norm(q_w, eps) -> linear + +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim]) + +-- kv_c (getitem 0) -> rms_norm(kv_w, eps) + +-- k_pe + """ + + def __init__( + self, + hidden_size: int, + q_dim: int = Q_DIM, + kv_c_dim: int = KV_C_DIM, + k_pe_dim: int = K_PE_DIM, + eps: float = EPS, + ): + super().__init__() + self.q_dim = q_dim + self.kv_dim = kv_c_dim + k_pe_dim + self.kv_c_dim = kv_c_dim + self.k_pe_dim = k_pe_dim + + self.proj = torch.nn.Linear(hidden_size, q_dim + self.kv_dim, bias=False) + self.q_norm = RMSNorm(q_dim, eps=eps) + self.kv_norm = RMSNorm(kv_c_dim, eps=eps) + self.q_b_proj = torch.nn.Linear(q_dim, hidden_size, bias=False) + + def forward( + self, x: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # Avoid graph input being a direct arg to a matched pattern node + x = torch.relu(x) + + projected = self.proj(x) + + q_c, kv_lora = projected.split([self.q_dim, self.kv_dim], dim=-1) + kv_c, k_pe = kv_lora.split([self.kv_c_dim, self.k_pe_dim], dim=-1) + + q_normed = self.q_norm(q_c) + kv_normed = self.kv_norm(kv_c) + + q_out = self.q_b_proj(q_normed) + return q_out, kv_normed, k_pe + + def ops_in_model_before(self): + return [torch.ops.vllm_ir.rms_norm.default] + + def ops_in_model_after(self): + return [torch.ops.vllm.fused_mla_dual_rms_norm.default] + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("hidden_size", [7168]) +@pytest.mark.skipif( + not is_aiter_found_and_supported(), + reason="Only test on ROCm with AITER installed and supported", +) +def test_fuse_mla_dual_rms_norm( + dtype: torch.dtype, + hidden_size: int, + monkeypatch: pytest.MonkeyPatch, +): + torch._dynamo.reset() + + vllm_config = VllmConfig( + model_config=ModelConfig(dtype=dtype), + compilation_config=CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + custom_ops=["+rms_norm"], + pass_config=PassConfig( + fuse_mla_dual_rms_norm=True, + eliminate_noops=True, + ), + ), + ) + + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m: + from vllm.compilation.passes.fusion.rocm_aiter_fusion import ( + MLADualRMSNormFusionPass, + ) + + torch.set_default_device("cuda") + torch.set_default_dtype(dtype) + torch.manual_seed(42) + + m.setenv("VLLM_ROCM_USE_AITER", "1") + rocm_aiter_ops.refresh_env_variables() + + fusion_pass = MLADualRMSNormFusionPass(vllm_config) + passes = [ + NoOpEliminationPass(vllm_config), + fusion_pass, + PostCleanupPass(vllm_config), + ] + backend = TestBackend(*passes) + model = MLADualRMSNormTestModel(hidden_size) + + x = torch.randn(1, hidden_size) + torch._dynamo.mark_dynamic(x, 0) + + outputs_unfused = model(x) + + model_fused = torch.compile(model, backend=backend) + outputs_fused = model_fused(x) + + torch.testing.assert_close(outputs_unfused, outputs_fused, atol=1e-2, rtol=1e-2) + + assert fusion_pass.matched_count == 1, ( + f"Expected 1 fused pair, got {fusion_pass.matched_count}" + ) + + backend.check_before_ops(model.ops_in_model_before()) + backend.check_after_ops(model.ops_in_model_after()) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 55d6d1297a8..650a229e177 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -962,6 +962,37 @@ def _rocm_aiter_triton_add_rmsnorm_pad_fake( return out, residual_out +def _fused_mla_dual_rms_norm_impl( + x1: torch.Tensor, + x1_weight: torch.Tensor, + x2: torch.Tensor, + x2_weight: torch.Tensor, + x1_epsilon: float, + x2_epsilon: float, +) -> tuple[torch.Tensor, torch.Tensor]: + from aiter.ops.fused_qk_norm_rope_cache_quant import fused_qk_rmsnorm + + return fused_qk_rmsnorm( + q=x1, + q_weight=x1_weight, + q_eps=x1_epsilon, + k=x2, + k_weight=x2_weight, + k_eps=x2_epsilon, + ) + + +def _fused_mla_dual_rms_norm_fake( + x1: torch.Tensor, + x1_weight: torch.Tensor, + x2: torch.Tensor, + x2_weight: torch.Tensor, + x1_epsilon: float, + x2_epsilon: float, +) -> tuple[torch.Tensor, torch.Tensor]: + return (torch.empty_like(x1), torch.empty_like(x2)) + + def _rocm_aiter_gemm_a8wfp4_impl( x: torch.Tensor, w: torch.Tensor, @@ -1491,6 +1522,13 @@ class rocm_aiter_ops: fake_impl=_triton_rotary_embedding_fake, ) + direct_register_custom_op( + op_name="fused_mla_dual_rms_norm", + op_func=_fused_mla_dual_rms_norm_impl, + mutates_args=[], + fake_impl=_fused_mla_dual_rms_norm_fake, + ) + _OPS_REGISTERED = True @staticmethod @@ -1537,6 +1575,10 @@ class rocm_aiter_ops: def get_triton_rotary_embedding_op() -> OpOverload: return torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default + @staticmethod + def get_fused_mla_dual_rms_norm_op() -> OpOverload: + return torch.ops.vllm.fused_mla_dual_rms_norm.default + @staticmethod def rms_norm( x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float diff --git a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py index 9a985472371..51b1a802f2e 100644 --- a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py +++ b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable from typing import Any import torch @@ -7,6 +8,7 @@ import torch._inductor.pattern_matcher as pm from torch import fx from torch._inductor.pattern_matcher import PatternMatcherPass +import vllm.ir.ops import vllm.model_executor.layers.quantization.utils.fp8_utils # noqa: F401 from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig @@ -20,7 +22,12 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.platforms import current_platform from ..inductor_pass import enable_fake_mode -from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass +from ..vllm_inductor_pass import ( + VllmFusionPatternMatcherPass, + VllmInductorPass, + VllmPatternMatcherPass, + VllmPatternReplacement, +) from .act_quant_fusion import ActivationQuantPattern from .matcher_utils import ( MatcherFusedAddRMSNorm, @@ -512,3 +519,101 @@ class RocmAiterTritonAddRMSNormPadFusionPass(VllmPatternMatcherPass): def uuid(self) -> str: return VllmInductorPass.hash_source(self, AddAiterRMSNormPadPattern) + + +class MLADualRMSNormPattern( + VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]] +): + """ + Fuse paired q_a_layernorm + kv_a_layernorm in MLA attention into + AITER's ``fused_qk_rmsnorm`` HIP kernel. + + Target FX-graph pattern (unfused, ``vllm_ir`` stage):: + + gemm -> split_with_sizes([q_dim, kv_dim]) + +-- q_c -> vllm_ir.rms_norm(q_c, q_w, eps) + +-- kv_lora -> split_with_sizes([kv_c_dim, k_pe_dim]) + +-- kv_c -> vllm_ir.rms_norm(kv_c, kv_w, eps) + +-- k_pe + + The pattern covers the connected subgraph rooted at the first + ``split_with_sizes`` (which produces ``q_c`` and ``kv_lora``), + through the two ``rms_norm`` calls, and the ``k_pe`` passthrough. + """ + + def __init__(self, epsilon: float) -> None: + self._epsilon = epsilon + + def get_inputs(self) -> list[torch.Tensor]: + q_dim, kv_c_dim, k_pe_dim = 8, 4, 2 + return [ + self.empty_bf16(5, q_dim + kv_c_dim + k_pe_dim), + self.empty_bf16(q_dim), + self.empty_bf16(kv_c_dim), + ] + + @property + def pattern( + self, + ) -> Callable[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + eps = self._epsilon + + def _pattern( + projected: torch.Tensor, + q_weight: torch.Tensor, + kv_weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q_dim = q_weight.shape[0] + kv_dim = projected.shape[-1] - q_dim + kv_c_dim = kv_weight.shape[0] + k_pe_dim = kv_dim - kv_c_dim + q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1) + kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1) + q_normed = vllm.ir.ops.rms_norm(q_c, q_weight, eps) + kv_normed = vllm.ir.ops.rms_norm(kv_c, kv_weight, eps) + return q_normed, kv_normed, k_pe + + return _pattern + + @property + def replacement( + self, + ) -> Callable[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + eps = self._epsilon + + def _replacement( + projected: torch.Tensor, + q_weight: torch.Tensor, + kv_weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q_dim = q_weight.shape[0] + kv_dim = projected.shape[-1] - q_dim + kv_c_dim = kv_weight.shape[0] + k_pe_dim = kv_dim - kv_c_dim + q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1) + kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1) + q_normed, kv_normed = torch.ops.vllm.fused_mla_dual_rms_norm( + q_c, + q_weight, + kv_c, + kv_weight, + eps, + eps, + ) + return q_normed, kv_normed, k_pe + + return _replacement + + +class MLADualRMSNormFusionPass(VllmFusionPatternMatcherPass): + """ + Post-grad PatternMatcher pass that fuses paired q / kv RMS norms in + MLA attention into ``fused_mla_dual_rms_norm`` backed by aiter's + ``fused_qk_rmsnorm`` HIP kernel. + """ + + def __init__(self, config: VllmConfig) -> None: + super().__init__(config, "mla_dual_rms_norm_fusion_pass") + + for epsilon in [1e-5, 1e-6]: + self.register(MLADualRMSNormPattern(epsilon)) diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index 91e10145607..b7c0d525c91 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -19,6 +19,7 @@ from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass if rocm_aiter_ops.is_enabled(): from .fusion.rocm_aiter_fusion import ( + MLADualRMSNormFusionPass, RocmAiterRMSNormQuantFusionPass, RocmAiterSiluMulFp8GroupQuantFusionPass, RocmAiterTritonAddRMSNormPadFusionPass, @@ -155,6 +156,9 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc] if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled(): self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)] + if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled(): + self.passes += [MLADualRMSNormFusionPass(config)] + if self.pass_config.fuse_rope_kvcache: self.passes += [SplitCoalescingPass(config)] self.passes += [ScatterSplitReplacementPass(config)] diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 6aca5c9825f..df86f579a1b 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -142,6 +142,8 @@ class PassConfig: # ROCm/AITER specific fusions fuse_act_padding: bool = None # type: ignore[assignment] """Fuse the custom RMSNorm + padding ops.""" + fuse_mla_dual_rms_norm: bool = None # type: ignore[assignment] + """Fuse paired q/kv RMS norms in MLA attention.""" fuse_rope_kvcache: bool = None # type: ignore[assignment] """Fuse the QK rope + KV cache ops.""" @@ -224,6 +226,7 @@ class PassConfig: "fuse_gemm_comms", "fuse_allreduce_rms", "fuse_act_padding", + "fuse_mla_dual_rms_norm", "fuse_rope_kvcache", mode="wrap", ) @@ -270,6 +273,12 @@ class PassConfig: "The fusion will be disabled." ) self.fuse_act_padding = False + if self.fuse_mla_dual_rms_norm and not current_platform.is_rocm(): + logger.warning_once( + "MLA dual RMS norm fusion requires ROCm/AITER. " + "The fusion will be disabled." + ) + self.fuse_mla_dual_rms_norm = False if self.fuse_rope_kvcache and not current_platform.is_rocm(): logger.warning_once( "KV cache fusion currently only enabled on ROCm. " diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 47bc3547ce8..7ca131b3631 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -165,6 +165,13 @@ def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool: ) +def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: + """Enable MLA dual RMS norm fusion when AITer is available.""" + from vllm._aiter_ops import rocm_aiter_ops + + return rocm_aiter_ops.is_enabled() + + OPTIMIZATION_LEVEL_00 = { "compilation_config": { "pass_config": { @@ -175,6 +182,7 @@ OPTIMIZATION_LEVEL_00 = { "enable_sp": False, "fuse_gemm_comms": False, "fuse_act_padding": False, + "fuse_mla_dual_rms_norm": False, "fuse_rope_kvcache": False, }, "cudagraph_mode": CUDAGraphMode.NONE, @@ -194,6 +202,7 @@ OPTIMIZATION_LEVEL_01 = { "enable_sp": False, "fuse_gemm_comms": False, "fuse_act_padding": enable_norm_pad_fusion, + "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": False, }, "cudagraph_mode": CUDAGraphMode.PIECEWISE, @@ -213,6 +222,7 @@ OPTIMIZATION_LEVEL_02 = { "enable_sp": IS_DENSE, "fuse_gemm_comms": IS_DENSE, "fuse_act_padding": enable_norm_pad_fusion, + "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, }, "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE, @@ -232,6 +242,7 @@ OPTIMIZATION_LEVEL_03 = { "enable_sp": IS_DENSE, "fuse_gemm_comms": IS_DENSE, "fuse_act_padding": enable_norm_pad_fusion, + "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, }, "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE, From 3a30eaa1d7b6497ce70c15558d410e37d8399b87 Mon Sep 17 00:00:00 2001 From: Hashem Hashemi <159079214+amd-hhashemi@users.noreply.github.com> Date: Mon, 20 Apr 2026 08:09:24 -0700 Subject: [PATCH 148/150] Properly enable wvSplitK fp8 path for RDNA (#37712) Signed-off-by: Hashem Hashemi --- vllm/model_executor/kernels/linear/scaled_mm/rocm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py index c8370dff512..64bc5b6c8bb 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py @@ -79,10 +79,10 @@ class ROCmFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import on_gfx12x, on_mi3xx - if not on_mi3xx(): - return False, "requires MI3xx." + if not (on_mi3xx() or on_gfx12x()): + return False, "requires MI3xx or gfx12x" if not envs.VLLM_ROCM_USE_SKINNY_GEMM: return False, "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled." From 595562651a5a4539ffa910d8570c08fb5169bdc9 Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Mon, 20 Apr 2026 23:31:39 +0800 Subject: [PATCH 149/150] [XPU] fix MoE triton backend in online fp8 quantization (#40109) Signed-off-by: Yan Ma --- vllm/model_executor/layers/fused_moe/oracle/fp8.py | 6 ++++++ vllm/model_executor/layers/fused_moe/xpu_fused_moe.py | 7 +++++++ vllm/model_executor/layers/quantization/fp8.py | 4 ---- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index d3c70d2d0a5..4420bb38731 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -471,6 +471,12 @@ def convert_to_fp8_moe_kernel_format( w2_input_scale=w2_input_scale, is_trtllm=(fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM), ) + elif fp8_backend == Fp8MoeBackend.XPU: + from vllm.model_executor.layers.fused_moe.xpu_fused_moe import ( + prepare_fp8_moe_layer_for_xpu, + ) + + w13, w2 = prepare_fp8_moe_layer_for_xpu(w13, w2) else: if fp8_backend not in [ Fp8MoeBackend.TRITON, diff --git a/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py index 9cc0ade288c..e10be4af868 100644 --- a/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/xpu_fused_moe.py @@ -24,6 +24,13 @@ if current_platform.is_xpu(): from vllm_xpu_kernels.fused_moe_interface import xpu_fused_moe +def prepare_fp8_moe_layer_for_xpu( + w13: torch.Tensor, + w2: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + return w13.transpose(-1, -2).contiguous(), w2.transpose(-1, -2).contiguous() + + class XPUExperts(mk.FusedMoEExpertsModular): def __init__( self, diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 4473bb8b2cc..0dc8907248e 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -1019,10 +1019,6 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod): layer.w2_weight[expert, :, :] ) - if current_platform.is_xpu(): - w13.data = w13.transpose(-1, -2).contiguous() - w2.data = w2.transpose(-1, -2).contiguous() - # Shuffle weights to runtime format and setup kernel. self._setup_kernel( layer, From 726efe177bf22874743d11dfdfef9247dbfb5ff0 Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Mon, 20 Apr 2026 12:28:46 -0400 Subject: [PATCH 150/150] [MoE Refactor] Move the shared/fused expert output sum into MoERunnerBase (#35949) Signed-off-by: Bill Nell Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- tests/kernels/moe/test_moe_layer.py | 93 ++----- .../test_shared_fused_moe_routed_transform.py | 36 +-- vllm/lora/layers/fused_moe.py | 3 - .../layers/fused_moe/fused_marlin_moe.py | 2 +- vllm/model_executor/layers/fused_moe/layer.py | 60 ++-- .../layers/fused_moe/modular_kernel.py | 2 +- .../fused_moe/runner/default_moe_runner.py | 4 - .../layers/fused_moe/runner/moe_runner.py | 13 +- .../fused_moe/runner/moe_runner_base.py | 261 ++++++++++++------ .../fused_moe/runner/moe_runner_factory.py | 6 +- .../layers/fused_moe/runner/shared_experts.py | 18 -- .../layers/fused_moe/shared_fused_moe.py | 8 +- vllm/model_executor/models/AXK1.py | 35 +-- vllm/model_executor/models/afmoe.py | 13 +- vllm/model_executor/models/aria.py | 8 +- vllm/model_executor/models/bailing_moe.py | 17 +- .../models/bailing_moe_linear.py | 19 +- vllm/model_executor/models/dbrx.py | 1 - vllm/model_executor/models/deepseek_v2.py | 34 +-- vllm/model_executor/models/dots1.py | 16 +- vllm/model_executor/models/ernie45_moe.py | 11 - vllm/model_executor/models/ernie45_vl_moe.py | 35 +-- vllm/model_executor/models/exaone_moe.py | 46 ++- vllm/model_executor/models/flex_olmo.py | 1 - vllm/model_executor/models/gemma4.py | 1 - vllm/model_executor/models/glm4_moe.py | 21 +- vllm/model_executor/models/gpt_oss.py | 1 - vllm/model_executor/models/granitemoe.py | 1 - vllm/model_executor/models/grok1.py | 1 - vllm/model_executor/models/hunyuan_v1.py | 7 - vllm/model_executor/models/interns1_pro.py | 1 - vllm/model_executor/models/jamba.py | 1 - vllm/model_executor/models/kimi_linear.py | 52 ++-- vllm/model_executor/models/lfm2_moe.py | 12 +- vllm/model_executor/models/llama4.py | 8 +- vllm/model_executor/models/longcat_flash.py | 1 - vllm/model_executor/models/mimo_v2_flash.py | 1 - vllm/model_executor/models/minimax_m2.py | 5 - vllm/model_executor/models/minimax_text_01.py | 1 - vllm/model_executor/models/mixtral.py | 1 - vllm/model_executor/models/nemotron_h.py | 29 +- vllm/model_executor/models/olmoe.py | 1 - vllm/model_executor/models/openpangu.py | 25 +- vllm/model_executor/models/param2moe.py | 16 +- vllm/model_executor/models/phimoe.py | 1 - vllm/model_executor/models/qwen2_moe.py | 7 - vllm/model_executor/models/qwen3_moe.py | 10 +- vllm/model_executor/models/qwen3_next.py | 8 - vllm/model_executor/models/sarvam.py | 16 +- vllm/model_executor/models/step3_text.py | 4 - vllm/model_executor/models/step3p5.py | 21 +- .../model_executor/models/transformers/moe.py | 30 +- 52 files changed, 325 insertions(+), 700 deletions(-) diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 01664d6a932..85619a91005 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -236,7 +236,6 @@ class MoETestConfig: use_gate: bool use_routed_input_transform: bool enable_eplb: bool = False - reduce_results: bool = False backend: str | None = None ep_size: int = 1 dp_size: int = 1 @@ -295,7 +294,6 @@ def generate_valid_test_configs( use_shared_experts, use_gate, use_routed_input_transform, - reduce_results, ) in product( SHAPE_COMBOS, NUM_EXPERTS, @@ -304,7 +302,6 @@ def generate_valid_test_configs( [False, True], # shared [False, True], # gate [False, True], # routed input exform - [False, True], # reduce results ): config = MoETestConfig( shape[0], # m @@ -318,7 +315,6 @@ def generate_valid_test_configs( use_gate, use_routed_input_transform, enable_eplb, - reduce_results, backend, ep_size, dp_size, @@ -395,18 +391,7 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: and config.backend.startswith("flashinfer_nvlink") and not current_platform.has_device_capability(90) ): - return False, "flashinfer_nvlink needs an H100+ GPUs" - - # reduce_results incompatibilities - if config.reduce_results and config.use_shared_experts: - return False, "reduce_results=True is not compatible with shared_experts=True" - - if config.reduce_results and config.quantization is not None: - return ( - False, - "reduce_results=True only tested with unquantized data types in " - "order to limit number of tests run", - ) + return False, "flashinfer_nvlink needs H100+ GPUs" # Backend-specific checks if config.backend is not None: @@ -448,10 +433,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: if config.enable_eplb and config.backend not in EPLB_SUPPORTED_BACKENDS: return False, f"EPLB not supported with {config.backend}." - world_size = config.tp_size * config.dp_size - if config.reduce_results and world_size == 1: - return False, "reduce_results=True only makes sense for multi-GPU tests" - if ( config.backend is not None and config.backend.startswith("flashinfer_nvlink") @@ -846,7 +827,6 @@ def make_fused_moe_layer( tp_size: int, ep_size: int, dp_size: int, - reduce_results: bool, w1: torch.Tensor, w2: torch.Tensor, top_k: int, @@ -874,7 +854,7 @@ def make_fused_moe_layer( routed_input_transform: torch.nn.Module | None = None, routed_output_transform: torch.nn.Module | None = None, pcp_size: int | None = 1, -) -> tuple[Callable, FusedMoE]: +) -> FusedMoE: quant_config, qw = make_quant_config(quantization, w1, w2, global_num_experts) kwargs = dict() @@ -887,8 +867,10 @@ def make_fused_moe_layer( # Add gate and routed_input_transform if provided if gate is not None: kwargs["gate"] = gate + if routed_input_transform is not None: kwargs["routed_input_transform"] = routed_input_transform + kwargs["routed_output_transform"] = routed_output_transform layer = builder( num_experts=global_num_experts, @@ -896,7 +878,6 @@ def make_fused_moe_layer( hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=in_dtype, - reduce_results=reduce_results, renormalize=renormalize, use_grouped_topk=use_grouped_topk, num_expert_group=num_expert_group, @@ -936,36 +917,7 @@ def make_fused_moe_layer( layer.quant_method.process_weights_after_loading(layer) - def _moe( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor: - if shared_experts is None: - final_shared_states = None - final_hidden_states = layer(hidden_states, router_logits) - else: - final_shared_states, final_hidden_states = layer( - hidden_states, router_logits - ) - - # Apply routed output transform if provided - # (e.g., latent space -> original space) - if routed_output_transform is not None: - final_hidden_states = routed_output_transform(final_hidden_states) - - if shared_experts is not None: - assert not reduce_results - assert final_shared_states is not None - final_hidden_states += final_shared_states - - if not reduce_results and layer.tp_size > 1: - final_hidden_states = layer.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - - return final_hidden_states - - return _moe, layer + return layer def make_fake_moe_layer( @@ -999,7 +951,6 @@ def make_fake_moe_layer( tp_size: int = 1, dp_size: int = 1, ep_size: int = 1, - reduce_results: bool = False, ) -> Callable: activation = MoEActivation.from_str(activation) @@ -1101,7 +1052,7 @@ def make_fake_moe_layer( def _test_body_regular( - moe_fn: Callable, + moe_layer: Callable, hidden_states: torch.Tensor, router_logits: torch.Tensor, vllm_config: VllmConfig, @@ -1118,13 +1069,12 @@ def _test_body_regular( num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ): - output = moe_fn(hidden_states, router_logits) + output = moe_layer(hidden_states, router_logits) return baseline_output, output def _test_body_eplb( - moe_fn: Callable, moe_layer: FusedMoE, hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -1145,7 +1095,6 @@ def _test_body_eplb( n: int, top_k: int, shared_experts, - reduce_results: bool, gate: torch.nn.Module | None, routed_input_transform: torch.nn.Module | None, routed_output_transform: torch.nn.Module | None, @@ -1161,7 +1110,7 @@ def _test_body_eplb( num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ): - output_before = moe_fn(hidden_states, router_logits) + output_before = moe_layer(hidden_states, router_logits) # Create a fresh FusedMoE layer with enable_eplb=True # Delete the original layer's registration so the constructor can @@ -1174,7 +1123,7 @@ def _test_body_eplb( # When using routed_input_transform, experts operate in latent space hidden_size_for_layer = k // 2 if routed_input_transform is not None else k - moe_fn, moe_layer = make_fused_moe_layer( + eplb_moe_layer = make_fused_moe_layer( quantization=quantization, use_ep=use_ep, hidden_size=hidden_size_for_layer, @@ -1183,7 +1132,6 @@ def _test_body_eplb( tp_size=tp_size, ep_size=ep_size, dp_size=dp_size, - reduce_results=reduce_results, w1=w1, w2=w2, top_k=top_k, @@ -1196,14 +1144,14 @@ def _test_body_eplb( ) # Necessary? - if moe_layer._expert_map is not None: - moe_layer._expert_map = moe_layer._expert_map.to(device) + if eplb_moe_layer._expert_map is not None: + eplb_moe_layer._expert_map = eplb_moe_layer._expert_map.to(device) # All ranks must generate the same permutation initial_indices = torch.arange(num_experts, dtype=torch.long) shuffled_indices = initial_indices[torch.randperm(num_experts)] - expert_weights = [list(moe_layer.get_expert_weights())] + expert_weights = [list(eplb_moe_layer.get_expert_weights())] communicator = create_eplb_communicator( group_coordinator=get_eplb_group(), @@ -1227,7 +1175,7 @@ def _test_body_eplb( num_experts, dtype=torch.int32, device=device ) - moe_layer.set_eplb_state( + eplb_moe_layer.set_eplb_state( moe_layer_idx=0, expert_load_view=torch.zeros( (1, num_experts), @@ -1244,7 +1192,7 @@ def _test_body_eplb( ), ) - moe_layer.eplb_state.should_record_tensor = torch.ones( + eplb_moe_layer.eplb_state.should_record_tensor = torch.ones( (), dtype=torch.bool, device=device ) @@ -1255,7 +1203,7 @@ def _test_body_eplb( num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ): - output_after = moe_fn(hidden_states, router_logits) + output_after = eplb_moe_layer(hidden_states, router_logits) return output_before, output_after @@ -1274,7 +1222,6 @@ def _run_one_config( num_experts: int, top_k: int, quantization: str | None, - reduce_results: bool, backend: str | None, test_body_fn: Callable, use_shared_experts: bool, @@ -1341,7 +1288,6 @@ def _run_one_config( tp_size=tp_size, ep_size=ep_size, dp_size=dp_size, - reduce_results=reduce_results, ) baseline_output = baseline_layer(hidden_states, router_logits) @@ -1369,7 +1315,7 @@ def _run_one_config( hidden_size_for_layer = k // 2 if routed_input_transform is not None else k # Create initial MoE layer - moe_fn, moe_layer = make_fused_moe_layer( + moe_layer = make_fused_moe_layer( quantization=quantization, use_ep=use_ep, hidden_size=hidden_size_for_layer, @@ -1378,7 +1324,6 @@ def _run_one_config( tp_size=tp_size, ep_size=ep_size, dp_size=dp_size, - reduce_results=reduce_results, w1=w1, w2=w2, top_k=top_k, @@ -1402,7 +1347,6 @@ def _run_one_config( # Call the test body function with all necessary context expected, actual = test_body_fn( - moe_fn=moe_fn, moe_layer=moe_layer, hidden_states=hidden_states, router_logits=router_logits, @@ -1423,7 +1367,6 @@ def _run_one_config( m=m, top_k=top_k, shared_experts=shared_experts, - reduce_results=reduce_results, gate=gate, routed_input_transform=routed_input_transform, routed_output_transform=routed_output_transform, @@ -1520,7 +1463,6 @@ def test_moe_layer_no_parallel( test_config.num_experts, test_config.top_k, test_config.quantization, - test_config.reduce_results, test_config.backend, _test_body_regular, use_shared_experts=test_config.use_shared_experts, @@ -1578,7 +1520,6 @@ def _parallel_worker( test_config.num_experts, test_config.top_k, test_config.quantization, - test_config.reduce_results, test_config.backend, functools.partial( _test_body_config, test_config=test_config, cpu_group=cpu_group @@ -1597,7 +1538,7 @@ def _parallel_worker( failed = failed + 1 if verbosity > 0: traceback.print_exc() - print(f"\n{str(ex)}\nFAILED {ex.__class__}") + print(f"\n{str(ex)}\nFAILED") else: print("F", end="") finally: diff --git a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py index 89e23eb0d74..464754c9f1b 100644 --- a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py +++ b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py @@ -165,7 +165,6 @@ def test_routed_input_transform_inside_vs_outside( top_k=top_k, hidden_size=latent_size, intermediate_size=intermediate_size, - reduce_results=False, renormalize=True, params_dtype=dtype, tp_size=1, @@ -183,7 +182,6 @@ def test_routed_input_transform_inside_vs_outside( top_k=top_k, hidden_size=latent_size, intermediate_size=intermediate_size, - reduce_results=False, renormalize=True, params_dtype=dtype, tp_size=1, @@ -212,34 +210,20 @@ def test_routed_input_transform_inside_vs_outside( hidden_states = torch.randn(num_tokens, hidden_size, device="cuda", dtype=dtype) router_logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype) - # Clone inputs so any in-place modification by Method A - # cannot affect Method B's computation. - hidden_states_A = hidden_states.clone() - router_logits_A = router_logits.clone() - with set_forward_context(None, vllm_config, num_tokens=num_tokens): - shared_out_A, routed_out_A = moe_with_transform( - hidden_states_A, router_logits_A - ) + # Method A: combined output (shared + routed) + combined_A = moe_with_transform(hidden_states, router_logits) + # Method B: manually transform, get routed output, add shared transformed_hidden = routed_transform(hidden_states) - shared_out_B, routed_out_B = moe_without_transform( - transformed_hidden, router_logits - ) + routed_out_B = moe_without_transform(transformed_hidden, router_logits) + shared_out_B = shared_experts(hidden_states) + combined_B = shared_out_B + routed_out_B - expected_shared_out = shared_experts(hidden_states) - - _assert_close( - routed_out_A, - routed_out_B, + torch.testing.assert_close( + combined_A, + combined_B, atol=1e-3, rtol=1e-3, - label="Routed output: transform inside vs outside", - ) - _assert_close( - shared_out_A, - expected_shared_out, - atol=1e-3, - rtol=1e-3, - label="Shared expert output", + msg="Combined output should match: transform inside vs outside", ) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 835bffe58ca..d6eec675c6d 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -592,9 +592,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): def forward(self, *args, **kwargs): return self.base_layer.forward(*args, **kwargs) - def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs): - return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs) - @property def quant_method(self): return self.base_layer.quant_method diff --git a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py index 6c916cf3cb6..daae5b6bd16 100644 --- a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py @@ -716,7 +716,7 @@ class MarlinExperts(MarlinExpertsBase): ): assert self.w1_scale is not None assert self.w2_scale is not None - return fused_marlin_moe( + fused_marlin_moe( hidden_states=hidden_states, w1=w1, w2=w2, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 190a9cc3b5d..bf10bc9d5c4 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -230,11 +230,18 @@ class FusedMoE(PluggableLayer): hidden_size: Input hidden state size of the transformer intermediate_size: Intermediate size of the experts params_dtype: Data type for the parameters. - reduce_results: Whether to all_reduce on the output of the layer renormalize: Whether to renormalize the logits in the fused_moe kernel quant_config: Quantization configure. enable_eplb: Whether to enable expert parallelism load balancer. router_logits_dtype: Data type for router logits buffers. + routed_scaling_factor: A scaling factor that is applied to the topk_weights + by the router or the output of the layer depending + on the value of `apply_routed_scale_to_output` + apply_routed_scale_to_output: Determine whether or not `routed_scaling_factor` + is applied to the topk_weights or to the experts + output. It is applied to the experts output + instead of the topk_weights when this feature is + not supported by the router (or the experts). """ # --8<-- [end:fused_moe] @@ -246,7 +253,6 @@ class FusedMoE(PluggableLayer): hidden_size: int, intermediate_size: int, params_dtype: torch.dtype | None = None, - reduce_results: bool = False, renormalize: bool = True, use_grouped_topk: bool = False, num_expert_group: int | None = None, @@ -274,12 +280,12 @@ class FusedMoE(PluggableLayer): gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, routed_input_transform: torch.nn.Module | None = None, + routed_output_transform: torch.nn.Module | None = None, + apply_routed_scale_to_output: bool = False, zero_expert_type: str | None = None, ): super().__init__() - self._routed_input_transform = routed_input_transform - if params_dtype is None: params_dtype = torch.get_default_dtype() self.params_dtype = params_dtype @@ -425,7 +431,6 @@ class FusedMoE(PluggableLayer): assert intermediate_size % self.tp_size == 0 intermediate_size_per_partition = intermediate_size // self.tp_size - self.reduce_results = reduce_results self.renormalize = renormalize # TODO(bnell): these attributes are only used by monolithic kernels. @@ -437,7 +442,14 @@ class FusedMoE(PluggableLayer): self.topk_group = topk_group self.custom_routing_function = custom_routing_function self.scoring_func = scoring_func - self.routed_scaling_factor = routed_scaling_factor + # When apply_routed_scale_to_output is True, we set the scaling factor + # to 1.0 so it ends up being a nop. Applying the scale will be handled + # by the runner in this case. + # The member variable must be set in the same way as the router since + # some quantization methods can access it. + self.routed_scaling_factor = ( + routed_scaling_factor if not apply_routed_scale_to_output else 1.0 + ) self.e_score_correction_bias = e_score_correction_bias # TODO(bnell): end attributes @@ -456,7 +468,7 @@ class FusedMoE(PluggableLayer): topk_group=topk_group, custom_routing_function=custom_routing_function, scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, + routed_scaling_factor=self.routed_scaling_factor, e_score_correction_bias=e_score_correction_bias, num_fused_shared_experts=self.num_fused_shared_experts, enable_eplb=enable_eplb, @@ -578,12 +590,18 @@ class FusedMoE(PluggableLayer): layer_name=self.layer_name, moe_config=self.moe_config, router=self.router, - routed_input_transform=self._routed_input_transform, gate=gate, shared_experts=shared_experts, quant_method=self.quant_method, - reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, + routed_input_transform=routed_input_transform, + routed_output_transform=routed_output_transform, + # When apply_routed_scale_to_output is True, we allow + # the scaling factor to be passed to the runner, otherwise + # we pass 1.0 so it ends up being a nop. + routed_scaling_factor=routed_scaling_factor + if apply_routed_scale_to_output + else 1.0, ) # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py @@ -1514,32 +1532,11 @@ class FusedMoE(PluggableLayer): self.ensure_moe_quant_config_init() return self.quant_method.moe_quant_config - def must_reduce_shared_expert_outputs(self) -> bool: - """ - The shared_experts are typically computed using the RowParallelLinear - layer. The result of this function is typically used as - the reduce_results argument to the module. - When just tensor-parallel is used, it is not required to reduce - the shared_experts results immediately. Instead we reduce at the - once at the end of the MoE op. (Refer to DeepSeekV2MoE module) - With EP and all2all kernels - this is no longer viable as all - GPU ranks in DP, produce the complete set of hidden_states. - Therefore it is required that we reduce the shared_experts output - early. - """ - return self.runner.must_reduce_shared_expert_outputs() - - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): - """ - Some combine kernels reduce across GPU ranks by default. - """ - return self.runner.maybe_all_reduce_tensor_model_parallel(final_hidden_states) - def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return self.runner.forward( hidden_states, router_logits, @@ -1613,7 +1610,6 @@ class FusedMoE(PluggableLayer): f"intermediate_size_per_partition={self.intermediate_size_per_partition}, " # noqa: E501 f"tp_size={self.tp_size},\n" f"ep_size={self.ep_size}, " - f"reduce_results={self.reduce_results}, " ) return s diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index f2e6e2560e7..376fcdf5b65 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1261,7 +1261,7 @@ class FusedMoEKernelModularImpl: topk_ids: torch.Tensor, apply_router_weight_on_input: bool, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: """ The _finalize method is a wrapper around self.prepare_finalize.finalize that handles DBO, async and shared expert overlap. diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 85c6563c084..8cd2fc65704 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -37,10 +37,6 @@ class DefaultMoERunner(MoERunnerBase): for different configurations (e.g., with/without shared experts, gates, etc.). """ - @property - def reduce_results(self) -> bool: - return self._reduce_results - @property def do_naive_dispatch_combine(self) -> bool: return ( diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 9ffbf3108f8..199ceab0659 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -26,18 +26,7 @@ class MoERunner(ABC): self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - raise NotImplementedError - - @abstractmethod - def must_reduce_shared_expert_outputs(self) -> bool: - raise NotImplementedError - - @abstractmethod - def maybe_all_reduce_tensor_model_parallel( - self, - final_hidden_states: torch.Tensor, - ): + ) -> torch.Tensor: raise NotImplementedError @abstractmethod diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 692d45d3460..2d2d0bb4ebe 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -81,7 +81,9 @@ def _resolve_layer_name(layer_name: str | LayerName) -> str: # Note: _moe_forward and _moe_forward_shared should not contain any # implementation details, They should merely pass along control to -# the runner's 'forward_dispatch' method. +# the runner's '_forward_dispatch' method. +# These functions should never be called directly since they do not +# include all the functionality of the MoE layer. def _moe_forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -89,7 +91,7 @@ def _moe_forward( layer_name: _layer_name_type, ) -> torch.Tensor: layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( + return layer.runner._forward_dispatch( layer, hidden_states, router_logits, @@ -113,7 +115,7 @@ def _moe_forward_shared( layer_name: _layer_name_type, ) -> tuple[torch.Tensor, torch.Tensor]: layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( + return layer.runner._forward_dispatch( layer, hidden_states, router_logits, @@ -143,7 +145,7 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward", op_func=_moe_forward, - mutates_args=["hidden_states"], # is this still true? + mutates_args=["hidden_states"], fake_impl=_moe_forward_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -157,6 +159,15 @@ direct_register_custom_op( ) +def _unpack( + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor | None, torch.Tensor]: + if isinstance(result, tuple): + return result + else: + return (None, result) + + class MoERunnerBase(MoERunner): """ Abstract base class providing common functionality for MoE runner implementations. @@ -174,7 +185,6 @@ class MoERunnerBase(MoERunner): allowing flexibility in the actual MoE computation implementation. Key abstract methods that subclasses must implement: - - reduce_results: Determines whether results should be reduced across ranks - _forward_impl: The core MoE computation logic specific to each runner type """ @@ -187,17 +197,23 @@ class MoERunnerBase(MoERunner): gate: torch.nn.Module | None, shared_experts: torch.nn.Module | None, quant_method: FusedMoEMethodBase, - reduce_results: bool, enable_dbo: bool, + routed_output_transform: torch.nn.Module | None = None, + routed_scaling_factor: float = 1.0, ): super().__init__() self.moe_config = moe_config self.router = router self.routed_input_transform = routed_input_transform + self.routed_output_transform = routed_output_transform + self.routed_scaling_factor = routed_scaling_factor self.gate = gate self.quant_method = quant_method - self._reduce_results = reduce_results self.enable_dbo = enable_dbo + self._fused_output_is_reduced = ( + self.quant_method.moe_kernel is not None + and self.quant_method.moe_kernel.output_is_reduced() + ) self._shared_experts: SharedExperts | None = None if shared_experts is not None: @@ -209,7 +225,6 @@ class MoERunnerBase(MoERunner): # called, i.e. by a MK or by the MoERunner. # Once the MK can be created upfront, we can just pass in the proper # flags derived from the quant_method's MK. - reduce_results=reduce_results, quant_method=quant_method, enable_dbo=enable_dbo, ) @@ -217,7 +232,7 @@ class MoERunnerBase(MoERunner): # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer_name - self.forward_entry = self._select_forward() + self._forward_entry = self._select_forward() def _select_forward(self) -> Callable: if current_platform.is_tpu() or current_platform.is_cpu(): @@ -245,38 +260,6 @@ class MoERunnerBase(MoERunner): def is_internal_router(self) -> bool: return self.gate is not None - @property - @abstractmethod - def reduce_results(self) -> bool: - raise NotImplementedError - - def must_reduce_shared_expert_outputs(self) -> bool: - """ - The shared_experts are typically computed using the RowParallelLinear - layer. The result of this function is typically used as - the reduce_results argument to the module. - When just tensor-parallel is used, it is not required to reduce - the shared_experts results immediately. Instead we reduce at the - once at the end of the MoE op. (Refer to DeepSeekV2MoE module) - With EP and all2all kernels - this is no longer viable as all - GPU ranks in DP, produce the complete set of hidden_states. - Therefore it is required that we reduce the shared_experts output - early. - """ - return ( - self.quant_method.moe_kernel is not None - and self.quant_method.moe_kernel.output_is_reduced() - ) - - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): - """ - Some combine kernels reduce across GPU ranks by default. - """ - if self.must_reduce_shared_expert_outputs(): - return final_hidden_states - else: - return tensor_model_parallel_all_reduce(final_hidden_states) - def apply_routed_input_transform( self, hidden_states: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor | None]: @@ -286,10 +269,6 @@ class MoERunnerBase(MoERunner): is saved separately so shared experts get [S, hidden_size] while routed experts get the transformed [S, moe_latent_size]. - TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be - moved inside SharedFusedMoE to all-reduce on the smaller latent - dimension. - Returns (possibly transformed) hidden states and the input for shared experts (or None if there are no shared experts). """ @@ -306,33 +285,79 @@ class MoERunnerBase(MoERunner): hidden_states if self._shared_experts is not None else None, ) - def _maybe_reduce_output( + def apply_routed_output_transform( self, - states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - trunc_sizes: list[int], - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return x[..., :trunc_size] + fused_output: torch.Tensor, + ) -> torch.Tensor: + """Apply transform to routed expert output (e.g., latent to full dim). - def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size) + Used by latent MoE models (e.g., NemotronH) where routed experts + operate in a compressed latent space and need projection back to + the full hidden dimension before combining with shared expert output. + """ + if self.routed_output_transform is not None: + r = self.routed_output_transform(fused_output) + fused_output = r[0] if isinstance(r, tuple) else r + return fused_output + def _maybe_apply_routed_scale_to_output( + self, + shared_output: torch.Tensor | None, + fused_output: torch.Tensor, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Apply routed_scaling_factor to the output with FP16 overflow + protection. + + Scale the fused expert output by routed_scaling_factor. For FP16, + avoid overflow by dividing shared_output by the scale instead + (the decoder layer compensates with matching divisions). + """ + if self.routed_scaling_factor != 1.0: + if fused_output.dtype != torch.float16: + fused_output *= self.routed_scaling_factor + elif shared_output is not None: + shared_output *= 1.0 / self.routed_scaling_factor + return shared_output, fused_output + + def _maybe_reduce_shared_expert_output( + self, + shared_output: torch.Tensor | None, + ) -> torch.Tensor | None: + """All-reduce shared expert output when the combine kernel already + reduced fused output. + + This is the "early" all-reduce path. When the combine kernel produces + already-reduced fused output, shared output must be reduced separately + to match. + """ + if self._fused_output_is_reduced: + assert shared_output is not None + shared_output = tensor_model_parallel_all_reduce(shared_output) + return shared_output + + def _maybe_reduce_final_output( + self, + states: torch.Tensor, + trunc_size: int, + ) -> torch.Tensor: + """Truncate padded dimensions and all-reduce the combined output. + + This is the "late" all-reduce path. When neither fused nor shared + output was individually reduced, the combined sum is all-reduced + here. Skipped when sequence-parallel is active (SP handles its + own reduction) or when the early path already reduced both outputs. + """ + # We don't need to reduce the final output if: + # - We are not running with TP or DP + # - The MK already reduced the fused output itself. if ( not self.moe_config.is_sequence_parallel - and self.reduce_results and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) + and not self._fused_output_is_reduced ): - func = reduce_and_trunc - else: - func = trunc + states = tensor_model_parallel_all_reduce(states) - if isinstance(states, tuple): - return tuple( - [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)] - ) - else: - assert len(trunc_sizes) == 1 - return func(states, trunc_sizes[0]) + return states[..., :trunc_size] def _encode_layer_name(self) -> str | LayerName: if _USE_LAYERNAME: @@ -349,7 +374,15 @@ class MoERunnerBase(MoERunner): self, shared_experts_input: torch.Tensor | None, hidden_states: torch.Tensor, - ) -> tuple[torch.Tensor, list[int]]: + ) -> tuple[torch.Tensor, int]: + """Pad hidden_states to moe_config.hidden_dim and compute the + original dimension for later truncation. + + For latent MoE, the routed hidden_states may be smaller than + hidden_dim. Padding ensures uniform tensor sizes through the + fused MoE kernel. The returned trunc_size is used by + _maybe_reduce_final_output to strip the padding from the result. + """ shared_experts_hidden_dim = ( shared_experts_input.shape[-1] if shared_experts_input is not None else 0 ) @@ -365,10 +398,10 @@ class MoERunnerBase(MoERunner): value=0.0, ) - if self._shared_experts is not None: - orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] + if self.routed_output_transform is not None and shared_experts_hidden_dim > 0: + orig_hidden_dims = shared_experts_hidden_dim else: - orig_hidden_dims = [transformed_hidden_dim] + orig_hidden_dims = transformed_hidden_dim return hidden_states, orig_hidden_dims @@ -388,6 +421,12 @@ class MoERunnerBase(MoERunner): router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Run expert routing and the fused MoE kernel via the quant method. + + Orchestrates shared expert execution (before/after), expert selection + via the router, and the actual fused MoE computation. Returns + (shared_expert_output, fused_expert_output). + """ # Run this before quant_method to avoid inplace issues. # TODO(bnell): probably not needed anymore since inplace is # disabled when shared experts are present. @@ -428,6 +467,13 @@ class MoERunnerBase(MoERunner): ) def _sequence_parallel_context(self): + """Return a context manager for sequence-parallel token + redistribution. + + When sequence parallelism is active, returns a context that handles + local size tracking for proper token scatter/gather. Otherwise + returns a no-op context. + """ ctx = get_forward_context() return ( ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) @@ -448,22 +494,25 @@ class MoERunnerBase(MoERunner): def _maybe_add_zero_expert_output( self, - result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + result: torch.Tensor, + ) -> torch.Tensor: + """Add the zero expert's contribution to the final result. + + When a ZeroExpertRouter is used, it computes a bias-like output + from the "zero expert" that is added to the combined routed+shared + expert output. + """ if isinstance(self.router, ZeroExpertRouter): zero_expert_output = self.router.zero_expert_output assert zero_expert_output is not None - if isinstance(result, tuple): - result = (result[0], result[1] + zero_expert_output) - else: - result = result + zero_expert_output + result = result + zero_expert_output return result def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: """Invoke the fused moe layer. Input: @@ -472,55 +521,88 @@ class MoERunnerBase(MoERunner): Output: - The new hidden_states. - or - - A tuple of (shared experts output, new hidden_states). Calling sequence - forward - - self.forward_entry (_moe_forward or _moe_forward_shared custom op) - - forward_dispatch + - self._forward_entry (_moe_forward or _moe_forward_shared custom op) + - _forward_dispatch - _forward_impl Note: The existence of _moe_forward and _moe_forward_shared custom ops are due to the following reasons: 1. the chunking loop in ChunkingMoERunner._forward_impl cannot be compiled by torch.compile - 2. pytorch cannot handle union types in custom op signatures so _moe_forward - and _moe_forward_shared must be split. + 2. pytorch cannot handle union types in custom op signatures so + _moe_forward and _moe_forward_shared must be split. If ChunkingMoERunner._forward_impl can be implemented via torch.scan we can potentially get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence into the 'forward' method. """ - # Apply transform for routed experts (e.g., latent projection for latent MoE) + # Apply transform for routed experts (e.g., latent projection + # for latent MoE) hidden_states, shared_experts_input = self.apply_routed_input_transform( hidden_states ) - hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( + hidden_states, og_hidden_dim = self._maybe_pad_hidden_states( shared_experts_input, hidden_states, ) - fused_output = self.forward_entry( + result = self._forward_entry( hidden_states, router_logits, shared_experts_input, self._encode_layer_name(), ) - result = self._maybe_reduce_output(fused_output, og_hidden_dims) + # + # Note: there are two all-reduce points below. They are mutually + # exclusive, controlled by _fused_output_is_reduced + # - When True: the combine kernel already reduced fused_output, + # so we reduce shared_output here to match, then skip the + # all-reduce in _maybe_reduce_final_output. + # - When False: neither output is reduced yet, so we combine + # them first and all-reduce the sum in _maybe_reduce_final_output. + + # Extract outputs from result + shared_output, fused_output = _unpack(result) + + # If combine kernel already reduced fused, reduce shared to match. + # See note above re: the two all-reduce points. + shared_output = self._maybe_reduce_shared_expert_output(shared_output) + + shared_output, fused_output = self._maybe_apply_routed_scale_to_output( + shared_output, fused_output + ) + + # Apply output transform (e.g. latent -> full dim) + fused_output = self.apply_routed_output_transform(fused_output) + + if shared_output is not None: + result = shared_output + fused_output + else: + result = fused_output + + result = self._maybe_reduce_final_output(result, og_hidden_dim) return self._maybe_add_zero_expert_output(result) - def forward_dispatch( + def _forward_dispatch( self, layer: torch.nn.Module, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Entry point called by the custom op to run the MoE computation. + + Handles pre-dispatch setup (gate application, external shared expert + triggering, quant config init) then delegates to _forward_impl within + the sequence-parallel context. + """ # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() @@ -549,4 +631,11 @@ class MoERunnerBase(MoERunner): router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Core MoE computation to be implemented by subclasses. + + Performs expert routing, fused MoE kernel execution, and shared + expert computation. Returns a single tensor (fused output only) + or a tuple of (shared_output, fused_output) when shared experts + are present. + """ raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py index 2143fa3ce08..feb4614d837 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py @@ -29,8 +29,9 @@ def create_moe_runner( gate: torch.nn.Module | None, shared_experts: SharedExperts | None, quant_method: FusedMoEMethodBase, - reduce_results: bool, enable_dbo: bool, + routed_output_transform: torch.nn.Module | None = None, + routed_scaling_factor: float = 1.0, ) -> MoERunner: return DefaultMoERunner( layer_name, @@ -40,6 +41,7 @@ def create_moe_runner( gate, shared_experts, quant_method, - reduce_results, enable_dbo, + routed_output_transform=routed_output_transform, + routed_scaling_factor=routed_scaling_factor, ) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 827a6e6bd3e..3a08d5d0fc2 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -5,10 +5,6 @@ from enum import IntEnum import torch import vllm.envs as envs -from vllm.distributed import ( - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, -) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -48,7 +44,6 @@ class SharedExperts: layer: torch.nn.Module, moe_config: FusedMoEConfig, quant_method: QuantizeMethodBase, - reduce_results: bool, enable_dbo: bool, ): from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( @@ -68,7 +63,6 @@ class SharedExperts: self._layer = layer self._moe_config = moe_config self._quant_method = quant_method - self._reduce_results = reduce_results # Allow disabling of the separate shared experts stream for # debug purposes. @@ -139,18 +133,6 @@ class SharedExperts: return output - def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: - # Reduce shared expert outputs if necessary, since the MLP - # should have been created with reduce_results=False. - if ( - self._reduce_results - and self._quant_method.moe_kernel is not None - and self._quant_method.moe_kernel.output_is_reduced() - and get_tensor_model_parallel_world_size() > 1 - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - return shared_out - @property def _output_idx(self) -> int: return dbo_current_ubatch_id() if self.enable_dbo else 0 diff --git a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py index ed243e99263..9cfcb1baa9b 100644 --- a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py @@ -18,12 +18,8 @@ class SharedFusedMoE(FusedMoE): self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - result = super().forward( + ) -> torch.Tensor: + return super().forward( hidden_states=hidden_states, router_logits=router_logits, ) - if self.shared_experts is None: - return None, result - else: - return result diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index f5ed4400fb6..d42fbed42ae 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -100,7 +100,7 @@ class AXK1MoE(nn.Module): self.tp_size = get_tensor_model_parallel_world_size() self.tp_rank = get_tensor_model_parallel_rank() - self.routed_scaling_factor = config.routed_scaling_factor + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) self.ep_group = get_ep_group().device_group self.ep_rank = get_ep_group().rank_in_group @@ -170,7 +170,6 @@ class AXK1MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -180,9 +179,8 @@ class AXK1MoE(nn.Module): scoring_func=config.scoring_func, # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally - routed_scaling_factor=1.0 - if not self.is_rocm_aiter_moe_enabled - else self.routed_scaling_factor, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -204,43 +202,20 @@ class AXK1MoE(nn.Module): hidden_states = sequence_parallel_chunk(hidden_states) if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) else: - # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.shared_experts is None: - assert shared_output is None - - # Fix FP16 overflow - # See AXK1DecoderLayer for more details. - if hidden_states.dtype != torch.float16: - if not self.is_rocm_aiter_moe_enabled: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - assert shared_output is not None - shared_output *= 1.0 / self.routed_scaling_factor - - if self.shared_experts is not None: - assert shared_output is not None - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/afmoe.py b/vllm/model_executor/models/afmoe.py index 22037336411..5bad52a0c49 100644 --- a/vllm/model_executor/models/afmoe.py +++ b/vllm/model_executor/models/afmoe.py @@ -131,7 +131,6 @@ class AfmoeMoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.route_norm if self.score_func == "sigmoid" else False, quant_config=quant_config, use_grouped_topk=True, @@ -152,20 +151,10 @@ class AfmoeMoE(nn.Module): router_logits = self.gate(hidden_states.to(dtype=torch.float32)) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_experts is not None: - shared_output, final_hidden_states = fused_moe_out - final_hidden_states = final_hidden_states + shared_output - else: - final_hidden_states = fused_moe_out - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py index 7b891f8ee42..7a079c56540 100644 --- a/vllm/model_executor/models/aria.py +++ b/vllm/model_executor/models/aria.py @@ -283,7 +283,6 @@ class AriaTextMoELayer(nn.Module): hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, quant_config=quant_config, - reduce_results=True, prefix=f"{prefix}.experts", ) @@ -301,12 +300,7 @@ class AriaTextMoELayer(nn.Module): router_output = torch.nn.functional.linear(hidden_states, self.router_weight) - sparse_expert_output = self.experts(hidden_states, router_output) - - if self.shared_experts is not None: - return sparse_expert_output[0] + sparse_expert_output[1] - else: - return sparse_expert_output + return self.experts(hidden_states, router_output) class AriaTextDecoderLayer(LlamaDecoderLayer): diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py index 7725dfa2a88..510d605f804 100644 --- a/vllm/model_executor/models/bailing_moe.py +++ b/vllm/model_executor/models/bailing_moe.py @@ -291,7 +291,6 @@ class BailingMoE(nn.Module): top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -301,6 +300,7 @@ class BailingMoE(nn.Module): topk_group=self.topk_group, use_grouped_topk=self.use_grouped_topk, router_logits_dtype=self.router_dtype, + routed_scaling_factor=self.routed_scaling_factor, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -314,21 +314,6 @@ class BailingMoE(nn.Module): final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - - if self.shared_experts is not None: - shared_output, final_hidden_states = final_hidden_states - else: - shared_output = None - - final_hidden_states *= self.routed_scaling_factor - - if shared_output is not None: - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_size) diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index ecc5d63ced7..a63ad83f45b 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -358,7 +358,6 @@ class BailingMoeV25(nn.Module): top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -368,6 +367,8 @@ class BailingMoeV25(nn.Module): topk_group=self.topk_group, use_grouped_topk=self.use_grouped_topk, router_logits_dtype=self.router_dtype, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -383,22 +384,6 @@ class BailingMoeV25(nn.Module): hidden_states=hidden_states, router_logits=router_logits ) - # Handle tuple return from SharedFusedMoE - if self.shared_experts is not None: - shared_output, final_hidden_states = final_hidden_states - else: - shared_output = None - - final_hidden_states *= self.routed_scaling_factor - - if shared_output is not None: - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(num_tokens, hidden_size) diff --git a/vllm/model_executor/models/dbrx.py b/vllm/model_executor/models/dbrx.py index ca6e6a49a98..a72f4e48716 100644 --- a/vllm/model_executor/models/dbrx.py +++ b/vllm/model_executor/models/dbrx.py @@ -85,7 +85,6 @@ class DbrxExperts(FusedMoE): hidden_size=config.d_model, intermediate_size=config.ffn_config.ffn_hidden_size, params_dtype=params_dtype, - reduce_results=True, renormalize=True, quant_config=quant_config, tp_size=get_tensor_model_parallel_world_size(), diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index cd28fb0192f..1b01caded94 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -318,7 +318,6 @@ class DeepseekV2MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -326,11 +325,9 @@ class DeepseekV2MoE(nn.Module): topk_group=getattr(config, "topk_group", 1), prefix=f"{prefix}.experts", scoring_func=getattr(config, "scoring_func", "softmax"), - # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally - routed_scaling_factor=1.0 - if not self.is_rocm_aiter_moe_enabled - else self.routed_scaling_factor, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -363,43 +360,20 @@ class DeepseekV2MoE(nn.Module): hidden_states = sequence_parallel_chunk(hidden_states) if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) else: - # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.shared_experts is None: - assert shared_output is None - - # Fix FP16 overflow - # See DeepseekV2DecoderLayer for more details. - if hidden_states.dtype != torch.float16: - if not self.is_rocm_aiter_moe_enabled: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - assert shared_output is not None - shared_output *= 1.0 / self.routed_scaling_factor - - if self.shared_experts is not None: - assert shared_output is not None - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/dots1.py b/vllm/model_executor/models/dots1.py index 4e393145462..c176b736568 100644 --- a/vllm/model_executor/models/dots1.py +++ b/vllm/model_executor/models/dots1.py @@ -37,7 +37,6 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention @@ -120,7 +119,6 @@ class Dots1MoE(nn.Module): prefix: str = "", ): super().__init__() - self.tp_size = get_tensor_model_parallel_world_size() self.routed_scaling_factor = config.routed_scaling_factor self.n_shared_experts = config.n_shared_experts @@ -163,7 +161,6 @@ class Dots1MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -171,9 +168,9 @@ class Dots1MoE(nn.Module): topk_group=config.topk_group, prefix=f"{prefix}.experts", scoring_func=config.scoring_func, - # we do scaling outside, set factor to 1.0 to avoid double mul - routed_scaling_factor=1.0, e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -182,16 +179,9 @@ class Dots1MoE(nn.Module): router_logits, _ = self.gate(hidden_states) - shared_out, routed_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_experts is not None: - final_hidden_states = (routed_out + shared_out) * self.routed_scaling_factor - else: - final_hidden_states = routed_out * self.routed_scaling_factor - - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py index f038cfb21f2..c92e230bcd2 100644 --- a/vllm/model_executor/models/ernie45_moe.py +++ b/vllm/model_executor/models/ernie45_moe.py @@ -194,7 +194,6 @@ class Ernie4_5_MoeMoE(nn.Module): top_k=config.moe_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -215,16 +214,6 @@ class Ernie4_5_MoeMoE(nn.Module): hidden_states=hidden_states, router_logits=router_logits ) - if self.has_shared_experts: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - else: - final_hidden_states = final_hidden_states[1] - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/ernie45_vl_moe.py b/vllm/model_executor/models/ernie45_vl_moe.py index 418fdcfa072..e4b7ac6fb00 100644 --- a/vllm/model_executor/models/ernie45_vl_moe.py +++ b/vllm/model_executor/models/ernie45_vl_moe.py @@ -263,7 +263,6 @@ class Ernie4_5_VLMoeMoE(nn.Module): top_k=config.moe_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size[0], - reduce_results=False, renormalize=True, quant_config=quant_config, e_score_correction_bias=self.e_score_correction_bias[0], @@ -301,7 +300,6 @@ class Ernie4_5_VLMoeMoE(nn.Module): top_k=config.moe_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size[1], - reduce_results=False, renormalize=True, quant_config=quant_config, e_score_correction_bias=self.e_score_correction_bias[1], @@ -342,9 +340,6 @@ class Ernie4_5_VLMoeMoE(nn.Module): visual_token_mask = visual_token_mask.repeat(1, self.hidden_size).bool() text_token_mask = ~visual_token_mask final_experts_hidden_states = torch.zeros_like(hidden_states) - final_shared_output = ( - torch.zeros_like(hidden_states) if self.has_shared_experts else None - ) text_hidden_states = hidden_states[text_token_mask].reshape( -1, self.hidden_size @@ -356,26 +351,20 @@ class Ernie4_5_VLMoeMoE(nn.Module): text_router_logits, _ = self.text_experts_gate( text_hidden_states.to(dtype=torch.float32) ) - text_shared_output, text_experts_output = self.text_experts( + text_output = self.text_experts( hidden_states=text_hidden_states, router_logits=text_router_logits ) - final_experts_hidden_states[text_token_mask] = text_experts_output.flatten() - if self.has_shared_experts: - final_shared_output[text_token_mask] = text_shared_output.flatten() + final_experts_hidden_states[text_token_mask] = text_output.flatten() vision_router_logits, _ = self.vision_experts_gate( vision_hidden_states.to(dtype=torch.float32) ) - vision_shared_output, vision_experts_output = self.vision_experts( + vision_output = self.vision_experts( hidden_states=vision_hidden_states, router_logits=vision_router_logits ) - final_experts_hidden_states[visual_token_mask] = ( - vision_experts_output.flatten() - ) - if self.has_shared_experts: - final_shared_output[visual_token_mask] = vision_shared_output.flatten() + final_experts_hidden_states[visual_token_mask] = vision_output.flatten() - final_hidden_states = (final_shared_output, final_experts_hidden_states) + final_hidden_states = final_experts_hidden_states else: # only text modal input text_router_logits, _ = self.text_experts_gate( @@ -386,20 +375,6 @@ class Ernie4_5_VLMoeMoE(nn.Module): hidden_states=hidden_states, router_logits=text_router_logits ) - if self.has_shared_experts: - # for shared_experts model - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - else: - # for not shared_experts model - final_hidden_states = final_hidden_states[1] - - if self.tp_size > 1: - final_hidden_states = ( - self.text_experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index d7282edcf4f..a46cadf007e 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -31,6 +31,7 @@ from vllm.distributed import ( get_tensor_model_parallel_world_size, ) from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -116,12 +117,26 @@ class ExaoneMoe(nn.Module): self.physical_expert_start + self.n_local_physical_experts ) - self.experts = FusedMoE( + if getattr(config, "num_shared_experts", 0) > 0: + intermediate_size = config.moe_intermediate_size * config.num_shared_experts + self.shared_experts = ExaoneMoeGatedMLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + else: + self.shared_experts = None + + self.experts = SharedFusedMoE( + shared_experts=self.shared_experts, + gate=self.gate, num_experts=self.n_routed_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -135,41 +150,16 @@ class ExaoneMoe(nn.Module): num_redundant_experts=self.n_redundant_experts, ) - if getattr(config, "num_shared_experts", 0) > 0: - intermediate_size = config.moe_intermediate_size * config.num_shared_experts - self.shared_experts = ExaoneMoeGatedMLP( - hidden_size=config.hidden_size, - intermediate_size=intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - reduce_results=self.experts.must_reduce_shared_expert_outputs(), - prefix=f"{prefix}.shared_experts", - ) - else: - self.shared_experts = None - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # NOTE: hidden_states can have either 1D or 2D shape. orig_shape = hidden_states.shape hidden_dim = hidden_states.shape[-1] hidden_states = hidden_states.view(-1, hidden_dim) - # router_logits: (num_tokens, n_experts) - router_logits, _ = self.gate(hidden_states) - final_hidden_states = self.experts( - hidden_states=hidden_states, router_logits=router_logits + hidden_states=hidden_states, router_logits=hidden_states ) - if self.shared_experts is not None: - shared_output = self.shared_experts(hidden_states) - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/flex_olmo.py b/vllm/model_executor/models/flex_olmo.py index 67be99a879f..1b2047eb231 100644 --- a/vllm/model_executor/models/flex_olmo.py +++ b/vllm/model_executor/models/flex_olmo.py @@ -76,7 +76,6 @@ class FlexOlmoMoE(nn.Module): top_k=hf_config.num_experts_per_tok, hidden_size=hf_config.hidden_size, intermediate_size=hf_config.intermediate_size, - reduce_results=True, renormalize=False, quant_config=None, tp_size=tp_size, diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index d166a9df38a..42762e36f81 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -349,7 +349,6 @@ class Gemma4MoE(nn.Module): "moe_intermediate_size", getattr(config, "expert_intermediate_size", None), ), - reduce_results=True, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index d0e6cb6ada8..671e868da0a 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -184,7 +184,6 @@ class Glm4MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -192,8 +191,8 @@ class Glm4MoE(nn.Module): topk_group=config.topk_group, prefix=f"{prefix}.experts", scoring_func="sigmoid", - # we do scaling outside, set factor to 1.0 to avoid double mul - routed_scaling_factor=1.0, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -207,23 +206,9 @@ class Glm4MoE(nn.Module): # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states.to(dtype=torch.float32)) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - - if self.shared_experts is not None: - shared_output, final_hidden_states = fused_moe_out - assert shared_output is not None - final_hidden_states = ( - final_hidden_states * self.routed_scaling_factor + shared_output - ) - else: - final_hidden_states = fused_moe_out * self.routed_scaling_factor - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 4e4eb581842..b6edc344302 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -189,7 +189,6 @@ class MLPBlock(torch.nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, - reduce_results=True, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/granitemoe.py b/vllm/model_executor/models/granitemoe.py index 171b2e0ec5a..f57a8c942bb 100644 --- a/vllm/model_executor/models/granitemoe.py +++ b/vllm/model_executor/models/granitemoe.py @@ -104,7 +104,6 @@ class GraniteMoeMoE(nn.Module): hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=True, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/grok1.py b/vllm/model_executor/models/grok1.py index 0bd6a8f3d60..c9aa3d2068f 100644 --- a/vllm/model_executor/models/grok1.py +++ b/vllm/model_executor/models/grok1.py @@ -209,7 +209,6 @@ class Grok1MoE(nn.Module): hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=renormalize, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/hunyuan_v1.py b/vllm/model_executor/models/hunyuan_v1.py index a0130402c66..35d30006a66 100644 --- a/vllm/model_executor/models/hunyuan_v1.py +++ b/vllm/model_executor/models/hunyuan_v1.py @@ -39,7 +39,6 @@ from vllm.distributed import ( get_ep_group, get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention @@ -445,7 +444,6 @@ class HunYuanSparseMoeBlock(nn.Module): top_k=top_k, hidden_size=config.hidden_size, intermediate_size=intermediate_size, - reduce_results=False, renormalize=top_k > 1, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -464,11 +462,6 @@ class HunYuanSparseMoeBlock(nn.Module): final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_mlp is not None: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 28331b8ef3e..9612ea57b2c 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -176,7 +176,6 @@ class InternS1ProMoeSparseMoeBlock(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=True, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/jamba.py b/vllm/model_executor/models/jamba.py index 980bcffb5f9..b4b3b6873db 100644 --- a/vllm/model_executor/models/jamba.py +++ b/vllm/model_executor/models/jamba.py @@ -90,7 +90,6 @@ class JambaMoE(nn.Module): self.intermediate_size, tp_size=tp_size, params_dtype=params_dtype, - reduce_results=True, renormalize=False, use_grouped_topk=False, quant_config=quant_config, diff --git a/vllm/model_executor/models/kimi_linear.py b/vllm/model_executor/models/kimi_linear.py index 4cd7b63c147..e586a3ac346 100644 --- a/vllm/model_executor/models/kimi_linear.py +++ b/vllm/model_executor/models/kimi_linear.py @@ -11,11 +11,10 @@ from vllm.config import CacheConfig, ModelConfig, ParallelConfig, VllmConfig from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import SharedFusedMoE from vllm.model_executor.layers.kda import KimiDeltaAttention from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -132,22 +131,6 @@ class KimiMoE(nn.Module): self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts)) - self.experts = FusedMoE( - num_experts=num_experts, - top_k=config.num_experts_per_token, - hidden_size=hidden_size, - intermediate_size=moe_intermediate_size, - reduce_results=False, - renormalize=moe_renormalize, - quant_config=quant_config, - use_grouped_topk=config.use_grouped_topk, - num_expert_group=config.num_expert_group, - topk_group=config.topk_group, - prefix=f"{prefix}.experts", - scoring_func=config.moe_router_activation_func, - e_score_correction_bias=self.gate.e_score_correction_bias, - ) - if self.num_shared_experts is not None: intermediate_size = moe_intermediate_size * self.num_shared_experts self.shared_experts = KimiMLP( @@ -158,22 +141,33 @@ class KimiMoE(nn.Module): reduce_results=False, prefix=f"{prefix}.shared_experts", ) + else: + self.shared_experts = None + + self.experts = SharedFusedMoE( + shared_experts=self.shared_experts, + num_experts=num_experts, + top_k=config.num_experts_per_token, + hidden_size=hidden_size, + intermediate_size=moe_intermediate_size, + renormalize=moe_renormalize, + quant_config=quant_config, + use_grouped_topk=config.use_grouped_topk, + num_expert_group=config.num_expert_group, + topk_group=config.topk_group, + prefix=f"{prefix}.experts", + scoring_func=config.moe_router_activation_func, + e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_size = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_size) - if self.num_shared_experts is not None: - shared_output = self.shared_experts(hidden_states) router_logits, _ = self.gate(hidden_states) - final_hidden_states = ( - self.experts(hidden_states=hidden_states, router_logits=router_logits) - * self.routed_scaling_factor + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits ) - if shared_output is not None: - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_size) @@ -482,7 +476,7 @@ class KimiLinearModel(nn.Module): if self.config.is_moe: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) - expert_params_mapping = FusedMoE.make_expert_params_mapping( + expert_params_mapping = SharedFusedMoE.make_expert_params_mapping( self, ckpt_gate_proj_name="w1", ckpt_down_proj_name="w2", diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index d955b7127ad..4b49430c1fa 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/vllm/model_executor/models/lfm2_moe.py @@ -150,7 +150,6 @@ class Lfm2MoeSparseMoeBlock(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, # needed for softmax score func @@ -161,6 +160,7 @@ class Lfm2MoeSparseMoeBlock(nn.Module): num_redundant_experts=self.n_redundant_experts, scoring_func="sigmoid", e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -170,16 +170,10 @@ class Lfm2MoeSparseMoeBlock(nn.Module): # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - final_hidden_states = ( - self.experts(hidden_states=hidden_states, router_logits=router_logits) - * self.routed_scaling_factor + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits ) - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index b84b4e2ae51..a1c0ac89605 100644 --- a/vllm/model_executor/models/llama4.py +++ b/vllm/model_executor/models/llama4.py @@ -135,7 +135,6 @@ class Llama4MoE(nn.Module): custom_routing_function=Llama4MoE.custom_routing_function, intermediate_size=intermediate_size_moe, apply_router_weight_on_input=True, - reduce_results=False, renormalize=False, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -151,19 +150,14 @@ class Llama4MoE(nn.Module): router_logits, _ = self.router(hidden_states) - shared_out, routed_out = self.experts( + experts_out = self.experts( hidden_states=hidden_states, router_logits=router_logits, ) - experts_out = routed_out + shared_out if self.is_sequence_parallel: experts_out = tensor_model_parallel_all_gather(experts_out, 0) experts_out = experts_out[:num_tokens] - elif self.tp_size > 1: - experts_out = self.experts.maybe_all_reduce_tensor_model_parallel( - experts_out - ) return experts_out diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py index 375b0b69b1f..945fcb61509 100644 --- a/vllm/model_executor/models/longcat_flash.py +++ b/vllm/model_executor/models/longcat_flash.py @@ -300,7 +300,6 @@ class LongcatMoe(nn.Module): top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, - reduce_results=True, params_dtype=params_dtype, renormalize=False, quant_config=quant_config, diff --git a/vllm/model_executor/models/mimo_v2_flash.py b/vllm/model_executor/models/mimo_v2_flash.py index 43475ed690c..0b466f16601 100644 --- a/vllm/model_executor/models/mimo_v2_flash.py +++ b/vllm/model_executor/models/mimo_v2_flash.py @@ -162,7 +162,6 @@ class MiMoV2MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=True, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py index 2a9e5f4e08a..84d8dda533f 100644 --- a/vllm/model_executor/models/minimax_m2.py +++ b/vllm/model_executor/models/minimax_m2.py @@ -36,7 +36,6 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import FusedMoE @@ -104,7 +103,6 @@ class MiniMaxM2MoE(nn.Module): e_score_correction_bias=self.e_score_correction_bias, hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, - reduce_results=False, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -134,9 +132,6 @@ class MiniMaxM2MoE(nn.Module): final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - final_hidden_states = final_hidden_states - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/minimax_text_01.py b/vllm/model_executor/models/minimax_text_01.py index 21d74d8b058..67d7cb2d8bc 100644 --- a/vllm/model_executor/models/minimax_text_01.py +++ b/vllm/model_executor/models/minimax_text_01.py @@ -162,7 +162,6 @@ class MiniMaxText01MoE(nn.Module): hidden_size=self.hidden_size, intermediate_size=self.intermediate_size * self.tp_size, params_dtype=self.params_dtype, - reduce_results=True, renormalize=True, quant_config=self.quant_config, tp_size=self.tp_size, diff --git a/vllm/model_executor/models/mixtral.py b/vllm/model_executor/models/mixtral.py index 376fd7a1709..c182444f667 100644 --- a/vllm/model_executor/models/mixtral.py +++ b/vllm/model_executor/models/mixtral.py @@ -132,7 +132,6 @@ class MixtralMoE(nn.Module): hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=True, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 8abbc808cce..fa068639648 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -216,7 +216,6 @@ class NemotronHMoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=self.moe_hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -231,6 +230,9 @@ class NemotronHMoE(nn.Module): num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, routed_input_transform=self.fc1_latent_proj, + routed_output_transform=self.fc2_latent_proj, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, router_logits_dtype=self.gate.out_dtype, ) @@ -244,38 +246,15 @@ class NemotronHMoE(nn.Module): # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - # SharedFusedMoE handles: - # - shared experts (with original hidden_states) - # - routed_input_transform (fc1_latent_proj) for latent MoE - # - multistream parallelism between shared and routed experts - shared_output, final_hidden_states = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - # Fix FP16 overflow - # See DeepseekV2DecoderLayer for more details. - if hidden_states.dtype != torch.float16: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - shared_output *= 1.0 / self.routed_scaling_factor - - # TODO: See SharedFusedMoE.apply_routed_input_transform - # for bandwidth optimization - if self.use_latent_moe: - final_hidden_states, _ = self.fc2_latent_proj(final_hidden_states) - - if self.shared_experts is not None: - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/olmoe.py b/vllm/model_executor/models/olmoe.py index f0afe0e997c..fcde2e41afb 100644 --- a/vllm/model_executor/models/olmoe.py +++ b/vllm/model_executor/models/olmoe.py @@ -98,7 +98,6 @@ class OlmoeMoE(nn.Module): top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, - reduce_results=True, renormalize=False, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 994ae82529a..7de84da5193 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -206,7 +206,6 @@ class OpenPanguMoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -214,8 +213,8 @@ class OpenPanguMoE(nn.Module): topk_group=1, prefix=f"{prefix}.experts", scoring_func="sigmoid", - # we do scaling outside, set factor to 1.0 to avoid double mul - routed_scaling_factor=1.0, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -234,33 +233,15 @@ class OpenPanguMoE(nn.Module): router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.shared_experts is None: - assert shared_output is None - - if hidden_states.dtype != torch.float16: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - assert shared_output is not None - shared_output *= 1.0 / self.routed_scaling_factor - - if self.shared_experts is not None: - assert shared_output is not None - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/param2moe.py b/vllm/model_executor/models/param2moe.py index 925f94e0600..fddd1a8f173 100644 --- a/vllm/model_executor/models/param2moe.py +++ b/vllm/model_executor/models/param2moe.py @@ -359,7 +359,6 @@ class Param2MoEMoEBlock(nn.Module): top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -388,24 +387,11 @@ class Param2MoEMoEBlock(nn.Module): self.gate.weight.float(), ).to(hidden_states.dtype) - final_hidden = self.experts( + expert_output = self.experts( hidden_states=hidden_states, router_logits=router_logits, ) - if self.shared_experts is not None: - shared_output, expert_output = final_hidden - else: - shared_output, expert_output = None, final_hidden - - if shared_output is not None: - expert_output = expert_output + shared_output - - if self.tp_size > 1: - expert_output = self.experts.maybe_all_reduce_tensor_model_parallel( - expert_output - ) - return expert_output.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/phimoe.py b/vllm/model_executor/models/phimoe.py index 0b55b7ec839..7d6083f202e 100644 --- a/vllm/model_executor/models/phimoe.py +++ b/vllm/model_executor/models/phimoe.py @@ -281,7 +281,6 @@ class PhiMoE(nn.Module): hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=False, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index 750835cccd9..b5d13e926d7 100644 --- a/vllm/model_executor/models/qwen2_moe.py +++ b/vllm/model_executor/models/qwen2_moe.py @@ -170,7 +170,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -187,12 +186,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module): final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_expert is not None: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index f2ce070be8b..f0f69d43537 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -212,7 +212,6 @@ class Qwen3MoeSparseMoeBlock(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -234,22 +233,15 @@ class Qwen3MoeSparseMoeBlock(nn.Module): # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - shared_out, fused_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - final_hidden_states = ( - shared_out + fused_out if shared_out is not None else fused_out - ) if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) # return to 1d if input is 1d return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 3bd026d9dc9..50d44dbbf63 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -153,7 +153,6 @@ class Qwen3NextSparseMoeBlock(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=getattr(config, "norm_topk_prob", True), quant_config=quant_config, prefix=f"{prefix}.experts", @@ -183,18 +182,11 @@ class Qwen3NextSparseMoeBlock(nn.Module): hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_expert is not None: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index fa5ec44d7e7..3656fc921b2 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -341,7 +341,6 @@ class SarvamMLAMoE(nn.Module): top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -370,20 +369,7 @@ class SarvamMLAMoE(nn.Module): router_logits=router_logits, ) - if self.shared_experts is not None: - shared_output, expert_output = final_hidden - else: - shared_output, expert_output = None, final_hidden - - if shared_output is not None: - expert_output = expert_output + shared_output - - if self.tp_size > 1: - expert_output = self.experts.maybe_all_reduce_tensor_model_parallel( - expert_output - ) - - return expert_output.view(num_tokens, hidden_dim) + return final_hidden.view(num_tokens, hidden_dim) class SarvamMLABlock(nn.Module): diff --git a/vllm/model_executor/models/step3_text.py b/vllm/model_executor/models/step3_text.py index 18b689166a5..636a121c590 100644 --- a/vllm/model_executor/models/step3_text.py +++ b/vllm/model_executor/models/step3_text.py @@ -14,7 +14,6 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul @@ -71,7 +70,6 @@ class FusedMoEBlock(nn.Module): top_k=config.moe_top_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_expert_weight, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -94,8 +92,6 @@ class FusedMoEBlock(nn.Module): final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index bb4bf14a963..018f7895602 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -379,7 +379,6 @@ class FusedMoEBlock(nn.Module): top_k=config.moe_top_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_expert_weight, quant_config=quant_config, activation=activation, @@ -397,30 +396,16 @@ class FusedMoEBlock(nn.Module): hidden_states = hidden_states.view(-1, hidden_dim) if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) else: - # router_logits: (num_tokens, n_experts) + # TODO(bnell): this gate could be moved into the FusedMoE? router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.share_expert is None: - assert shared_output is None - - if self.share_expert is not None: - assert shared_output is not None - final_hidden_states += shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 81d21abbd06..cf13958ef76 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -204,8 +204,6 @@ class MoEMixin(MixtureOfExperts): ) assert intermediate_size is not None - # If there are shared experts, the results are - # reduced after mlp.forward() not inside FusedMoE num_shared_experts = getattr_iter( text_config, [ @@ -214,17 +212,6 @@ class MoEMixin(MixtureOfExperts): ], 0, ) - reduce_results = num_shared_experts == 0 - - def add_all_reduce(mlp: nn.Module): - """Adds an all-reduce to the output of `mlp.forward()`.""" - - class MLPWithAllReduce(mlp.__class__): - def forward(self, *args, **kwargs): - output = super().forward(*args, **kwargs) - return self.experts.maybe_all_reduce_tensor_model_parallel(output) - - mlp.__class__ = MLPWithAllReduce # Unused kwargs since we use custom_routing_function: # - `scoring_func` and `e_score_correction_bias` only used for grouped @@ -289,14 +276,11 @@ class MoEMixin(MixtureOfExperts): if "bias" in experts_param_name: has_bias = True break - # Double check there are no shared experts - nonlocal reduce_results - if reduce_results: + # If the config does not specify num_shared_experts, but + # the model has shared experts, we assume there is one. + if self.num_shared_experts == 0: for mlp_param_name, _ in mlp.named_parameters(): if "shared_expert" in mlp_param_name: - reduce_results = False - # If the config does not specify num_shared_experts, but - # the model has shared experts, we assume there is one. self.num_shared_experts = 1 break # Replace experts module with FusedMoE @@ -305,7 +289,6 @@ class MoEMixin(MixtureOfExperts): top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, - reduce_results=reduce_results, renormalize=renormalize, # Hard coded because topk happens in Transformers use_grouped_topk=False, @@ -326,13 +309,6 @@ class MoEMixin(MixtureOfExperts): self.moe_layers.append(fused_experts) self.expert_weights.append(fused_experts.get_expert_weights()) self.num_moe_layers += 1 - # If results are not all-reduced in FusedMoE, ensure they - # are all-reduced at the end of mlp.forward() if tensor - # parallel or expert parallel is enabled - if not reduce_results and ( - fused_experts.tp_size > 1 or fused_experts.ep_size > 1 - ): - add_all_reduce(mlp) else: _recursive_replace(child_module, prefix=qual_name)